import pyxel
import random
import json
import os
from collections import deque

SAVE_FILE = "deepcore_save.json"

def load_best():
    if os.path.exists(SAVE_FILE):
        try:
            with open(SAVE_FILE) as f:
                return json.load(f).get("best_blocks", None)
        except Exception:
            pass
    return None

def save_best(blocks):
    try:
        with open(SAVE_FILE, "w") as f:
            json.dump({"best_blocks": blocks}, f)
    except Exception:
        pass

TILE     = 8
W        = 40
H        = 30
SCREEN_W = W * TILE
SCREEN_H = H * TILE
VISION_RADIUS = 4
MAX_OXYGEN    = 900

# Ranks: (min_blocks, max_blocks, label, color)
RANKS = [
    (0,   40,  "DEEPCORE LEGEND",  10),
    (41,  70,  "Master Miner",      9),
    (71,  110, "Expert Digger",    11),
    (111, 160, "Seasoned Miner",    7),
    (161, 220, "Apprentice",       13),
    (221, 300, "Rock Cracker",      6),
    (301, 9999,"Tunnel Rat",        4),
]

def get_rank(blocks):
    for lo, hi, label, col in RANKS:
        if lo <= blocks <= hi:
            return label, col
    return "Tunnel Rat", 4

# Ore weights: higher numbers = harder rock type (1=dirt, 2=stone, 3=granite)
# Harder distribution than before — more stone and granite throughout
LEVELS = [
    {"name":"Surface",  "bg":1, "title_col":7,  "ore_weights":[1,2,2,2,2,3], "drain":2, "music_ch":0},
    {"name":"Midcore",  "bg":1, "title_col":9,  "ore_weights":[1,2,2,3,3,3], "drain":2, "music_ch":1},
    {"name":"Deepcore", "bg":1, "title_col":8,  "ore_weights":[2,2,3,3,3,3], "drain":1, "music_ch":2},
]

ROCK = {
    0: {"name":"empty",  "hard":0, "color":0},
    1: {"name":"dirt",   "hard":1, "color":4},
    2: {"name":"stone",  "hard":2, "color":5},
    3: {"name":"granite","hard":3, "color":6},
}

PICKAXE_COLOR = [10, 9]

# ── Particles ──────────────────────────────────────────────────────────────
class Particle:
    def __init__(self, x, y, col):
        self.x    = float(x)
        self.y    = float(y)
        self.vx   = random.uniform(-1.2, 1.2)
        self.vy   = random.uniform(-2.0, -0.3)
        self.col  = col
        self.life = random.randint(8, 18)
    def update(self):
        self.x += self.vx; self.vy += 0.18; self.y += self.vy; self.life -= 1
    def draw(self):
        if self.life > 0:
            px, py = int(self.x), int(self.y)
            if 0 <= px < SCREEN_W and 0 <= py < SCREEN_H:
                pyxel.pset(px, py, self.col)

# ── World generation ───────────────────────────────────────────────────────
# Make Sure p1, p2 are reachable
def bfs_reachable(world, sx, sy, pick):
    seen = {(sx, sy)}; q = deque([(sx, sy)])
    while q:
        x, y = q.popleft()
        for dx, dy in ((-1,0),(1,0),(0,-1),(0,1)):
            nx, ny = x+dx, y+dy
            if 0 <= nx < W and 0 <= ny < H and (nx,ny) not in seen:
                if ROCK[world[ny][nx]]["hard"] <= pick:
                    seen.add((nx,ny)); q.append((nx,ny))
    return seen

def _carve(world, x1, y1, x2, y2, mh):
    cx, cy = x1, y1
    while cx != x2:
        cx += 1 if x2 > cx else -1
        if ROCK[world[cy][cx]]["hard"] > mh: world[cy][cx] = mh if mh else 0
    while cy != y2:
        cy += 1 if y2 > cy else -1
        if ROCK[world[cy][cx]]["hard"] > mh: world[cy][cx] = mh if mh else 0

def carve_tunnel(world, x1, y1, x2, y2, mh):
    n = random.randint(3,5); pts = [(x1,y1)]
    for i in range(n):
        wx = random.randint(2, W-3)
        wy = random.randint(max(2, y1+(y2-y1)*i//(n+1)-2), min(H-2, y1+(y2-y1)*(i+1)//(n+1)+2))
        pts.append((wx,wy))
    pts.append((x2,y2))
    for i in range(len(pts)-1):
        ax,ay = pts[i]; bx,by = pts[i+1]
        if random.random() < 0.5: _carve(world,ax,ay,bx,by,mh)
        else: _carve(world,ax,ay,ax,by,mh); _carve(world,ax,by,bx,by,mh)

def carve_branch_tunnels(world, sx, sy, p2x, p2y, p3x, p3y):
    """Add extra branch tunnels so there are always multiple paths to explore."""
    n_branches = random.randint(2, 4)
    anchors = [
        (sx, sy),
        (p2x, p2y),
        (p3x, p3y),
    ]
    for _ in range(n_branches):
        ax, ay = random.choice(anchors)
        bx = random.randint(2, W-3)
        by = random.randint(max(3, ay-8), min(H-3, ay+8))
        carve_tunnel(world, ax, ay, bx, by, 1)
        if random.random() < 0.5:
            sx2 = bx + random.randint(-6, 6)
            sy2 = by + random.randint(-4, 4)
            sx2 = max(2, min(W-3, sx2))
            sy2 = max(3, min(H-3, sy2))
            carve_tunnel(world, bx, by, sx2, sy2, 1)

# (Ehrenlos) 
def add_dead_ends(world, count=10):
    for _ in range(count):
        sx,sy = random.randint(1,W-2), random.randint(4,H-4)
        dx,dy = random.choice([(1,0),(-1,0),(0,1),(0,-1)])
        for i in range(random.randint(4,10)):
            nx,ny = sx+dx*i, sy+dy*i
            if 0<=nx=n: break
            tx,ty=random.randint(1,W-2),random.randint(4,H-2)
            if (tx,ty) in safe: continue
            if empty and world[ty][tx]!=0: continue
            if not empty and world[ty][tx]==0: continue
            col.append((tx,ty))
    place(gas, n_g, False)
    cands=[]; place(cands, n_sp*3, True)
    sp_set=set()
    for cx,cy in cands:
        sp_set.add((cx,cy))
        if (p2x,p2y) not in _bfs_spikes(world,sx,sy,1,sp_set): sp_set.discard((cx,cy))
        if len(sp_set)>=n_sp: break
    return list(sp_set), gas

def generate_world(sx, sy, lvl):
    w=LEVELS[lvl]["ore_weights"]
    world=[[0 if y<3 else random.choice(w) for x in range(W)] for y in range(H)]
    p2x=random.randint(4,W-5); p2y=random.randint(7,H//2)
    p3x=random.randint(4,W-5); p3y=random.randint(H//2+3,H-4)
    while abs(p3x-p2x)<4 and abs(p3y-p2y)<4:
        p3x=random.randint(4,W-5); p3y=random.randint(H//2+3,H-4)
    world[p2y][p2x]=0; world[p3y][p3x]=0
    carve_tunnel(world,sx,sy,p2x,p2y,1)
    carve_tunnel(world,p2x,p2y,p3x,p3y,2)
    # Add extra branch tunnels for multiple exploration paths
    carve_branch_tunnels(world, sx, sy, p2x, p2y, p3x, p3y)
    add_dead_ends(world)
    if (p2x,p2y) not in bfs_reachable(world,sx,sy,1): carve_tunnel(world,sx,sy,p2x,p2y,0)
    if (p3x,p3y) not in bfs_reachable(world,sx,sy,2): carve_tunnel(world,p2x,p2y,p3x,p3y,0)
    tanks=add_o2_tanks(world)
    picks=[(p2x,p2y,2),(p3x,p3y,3)]
    ex,ey=p3x,min(p3y+3,H-2); world[ey][ex]=0
    safe=set()
    for px2,py2 in [(sx,sy),(p2x,p2y),(p3x,p3y),(ex,ey)]:
        for dy in range(-2,3):
            for dx in range(-2,3): safe.add((px2+dx,py2+dy))
    for tx,ty in tanks: safe.add((tx,ty))
    spikes,gas=add_traps(world,lvl,safe,sx,sy,p2x,p2y)
    return world,picks,tanks,(ex,ey),spikes,gas

# ── Sound ──────────────────────────────────────────────────────────────────
def _snd(i, notes, spd=10):
    n=len(notes); pyxel.sounds[i].set("".join(notes),"t"*n,"5"*n,"n"*n,spd)

def setup_sounds():
    _snd(0,["c2"],spd=6); _snd(1,["g1"],spd=6)
    _snd(2,["c3","e3","g3","c4"],spd=8)
    _snd(3,["c4","e4","g4","c4"],spd=8)
    _snd(4,["c2","a1","f1","d1","c1"],spd=10)
    _snd(5,["c3","e3","g3","c4","e4","g4","c4"],spd=8)
    _snd(6,["g2","c3","e3","g3"],spd=8)
    _snd(7,["a3","a3","a3"],spd=4); _snd(8,["e3","e3","e3"],spd=4)
    _snd(10,["c3","c3","e3","g3","e3","c3","d3","f3","a3","f3","d3","b2","c3","e3","g3","c4"],spd=16)
    _snd(11,["c2","c2","g2","g2","a2","a2","f2","f2","c2","c2","g2","g2","f2","f2","c2","c2"],spd=16)
    pyxel.music(0).set([10],[11],[],[])
    _snd(20,["a2","b2","c3","b2","a2","g2","f2","g2","a2","b2","c3","d3","e3","d3","c3","b2"],spd=14)
    _snd(21,["a1","a1","e2","e2","f2","f2","d2","d2","a1","a1","e2","e2","d2","d2","a1","a1"],spd=14)
    pyxel.music(1).set([20],[21],[],[])
    _snd(30,["f2","f2","e2","f2","g2","f2","e2","d2","c2","d2","e2","f2","g2","a2","b2","c3"],spd=12)
    _snd(31,["f1","f1","c2","c2","g1","g1","d2","d2","f1","f1","c2","c2","d2","d2","f1","f1"],spd=12)
    pyxel.music(2).set([30],[31],[],[])

# ── Tile drawing ───────────────────────────────────────────────────────────
def draw_tile(x, y, block):
    bx, by = x*TILE, y*TILE
    if block == 1:
        pyxel.rect(bx,by,TILE,TILE,4)
        for ox,oy in [(1,1),(5,2),(2,5),(6,4),(3,3),(7,6),(0,6),(4,0)]: pyxel.pset(bx+ox,by+oy,2)
        for ox,oy in [(3,1),(6,3),(1,6),(5,5)]: pyxel.pset(bx+ox,by+oy,9)
    elif block == 2:
        pyxel.rect(bx,by,TILE,TILE,13)
        pyxel.rect(bx+1,by+1,3,2,5); pyxel.rect(bx+4,by+4,3,2,6)
        pyxel.pset(bx+6,by+1,5); pyxel.pset(bx+1,by+5,6)
        pyxel.line(bx,by+3,bx+3,by+3,4); pyxel.line(bx+5,by+6,bx+7,by+6,4)
    elif block == 3:
        pyxel.rect(bx,by,TILE,TILE,6)
        for ox,oy,c in [(1,2,13),(5,1,13),(3,5,13),(6,4,7),(2,4,5),(4,2,5),(6,6,1),(1,6,5)]:
            pyxel.pset(bx+ox,by+oy,c)
        pyxel.line(bx+2,by+1,bx+3,by+3,5); pyxel.line(bx+5,by+5,bx+6,by+7,1)

def draw_player(px, py, f):
    bx,by=px*TILE,py*TILE
    pyxel.rect(bx+2,by+3,4,4,11); pyxel.rect(bx+2,by+1,4,3,7)
    pyxel.rect(bx+1,by,6,2,9); pyxel.pset(bx+3,by,10)
    if (f//6)%2==0: pyxel.rect(bx+2,by+7,2,1,7); pyxel.rect(bx+4,by+6,2,1,7)
    else:           pyxel.rect(bx+2,by+6,2,1,7); pyxel.rect(bx+4,by+7,2,1,7)

# ── Guardian ───────────────────────────────────────────────────────────────
class Guardian:
    def __init__(self, x, y, interval, ca, cb):
        self.x=x; self.y=y; self.interval=interval; self.ca=ca; self.cb=cb; self._t=0
    def update(self, world, px, py):
        self._t+=1
        if self._t < self.interval: return False
        self._t=0
        path=self._bfs(world,px,py)
        if path and len(path)>1: self.x,self.y=path[1]
        return self.x==px and self.y==py
    def _bfs(self, world, tx, ty):
        s=(self.x,self.y); g=(tx,ty)
        if s==g: return [s]
        vis={s:None}; q=deque([s])
        while q:
            cx,cy=q.popleft()
            if (cx,cy)==g:
                path=[]; cur=g
                while cur: path.append(cur); cur=vis[cur]
                path.reverse(); return path
            for dx,dy in ((-1,0),(1,0),(0,-1),(0,1)):
                nx,ny=cx+dx,cy+dy
                if (nx,ny) not in vis and 0<=nx=h:
            self._sparks(nx,ny,ROCK[blk]["color"])
            pyxel.play(3,0 if h==1 else 1)
            self.world[ny][nx]=0; self.px,self.py=nx,ny
            self.total_score+=h*10; self.blocks_mined+=1
        else: self.disc[ny][nx]=True; return
        self._vis(); self._pickup(); self._o2tank(); self._traps(nx,ny)

    def _traps(self,x,y):
        if (x,y) in self.spikes:
            self.msg="Impaled on spikes!"; self.mt=120; self.why="spike"; self._die()

    def _die(self, why=None):
        if why: self.why=why
        self.dead=True; pyxel.stop(); pyxel.play(3,4)

    def _sparks(self,tx,ty,col):
        cx,cy=tx*TILE+TILE//2,ty*TILE+TILE//2
        for _ in range(random.randint(4,8)): self.parts.append(Particle(cx,cy,col))

    def _pickup(self):
        for item in self.picks[:]:
            ix,iy,lv=item
            if self.px==ix and self.py==iy:
                self.picks.remove(item)
                if lv>self.pickaxe: self.pickaxe=lv; self.msg=f"Pickaxe Lv.{lv}! Upgrade!"
                else: self.msg=f"Pickaxe Lv.{lv} (have better)"
                self.mt=150; pyxel.play(3,2); self.total_score+=200

    def _o2tank(self):
        for t in self.tanks[:]:
            tx,ty=t
            if self.px==tx and self.py==ty:
                self.tanks.remove(t); self.o2=MAX_OXYGEN
                self.msg="O2 refilled!"; self.mt=90; pyxel.play(3,3); self.total_score+=100

    def _exit(self):
        ex,ey=self.exit
        if self.px==ex and self.py==ey and not self.picks:
            if self.level_idx=self.tt: self.level_idx+=1; self.reset_level(new_game=False)
            return
        if self.dead or self.won:
            if pyxel.btnp(pyxel.KEY_R) or pyxel.btnp(pyxel.KEY_RETURN):
                self.game_state="menu"; pyxel.stop()
            return
        drain=LEVELS[self.level_idx]["drain"]
        if self.py>=3:
            if pyxel.frame_count%drain==0: self.o2-=1
            for gx,gy in self.gas:
                if abs(self.px-gx)<=2 and abs(self.py-gy)<=2:
                    if pyxel.frame_count%2==0: self.o2-=1
            if self.o2<=0: self.o2=0; self.why="oxygen"; self._die(); return
        if self.o20]
        if self.mt>0: self.mt-=1
        self._exit()

    # ── Menu ───────────────────────────────────────────────────────────────
    def _menu(self):
        f=self._mf; pyxel.cls(1)
        for i in range(35):
            pyxel.rect((i*19+f//2)%SCREEN_W,(i*13+f)%SCREEN_H,2,2,[4,2,6,5,13][i%5])
        # Title
        pyxel.rect(SCREEN_W//2-90,8,180,20,0); pyxel.rectb(SCREEN_W//2-90,8,180,20,6)
        pyxel.text(SCREEN_W//2-60,14,"D E E P C O R E   M I N E R",13)
        pyxel.text(SCREEN_W//2-54,32,"Dig deep.  Breathe fast.  Escape.",9 if (f//20)%2==0 else 10)
        # Characters
        walk=(f//10)%2; mx,my=SCREEN_W//2-38,44
        pyxel.rect(mx+2,my+3,4,4,11); pyxel.rect(mx+2,my+1,4,3,7)
        pyxel.rect(mx+1,my,6,2,9); pyxel.pset(mx+3,my,10)
        pyxel.rect(mx+2,my+7-walk,2,1,7); pyxel.rect(mx+4,my+6+walk,2,1,7)
        pyxel.text(mx+1,my+10,"YOU",11)
        pyxel.text(SCREEN_W//2-4,my+4,"VS",7)
        gc1=8 if (f//8)%2==0 else 14; gx1,gy1=SCREEN_W//2+12,44
        pyxel.rect(gx1+2,gy1+2,4,5,gc1); pyxel.pset(gx1+3,gy1+3,7); pyxel.pset(gx1+5,gy1+3,7)
        pyxel.rect(gx1+2,gy1+7-walk,2,1,gc1); pyxel.rect(gx1+5,gy1+6+walk,1,1,gc1)
        pyxel.text(gx1-2,gy1+10,"G1",gc1)
        gc2=2 if (f//10)%2==0 else 14; gx2,gy2=SCREEN_W//2+28,44
        pyxel.rect(gx2+2,gy2+2,4,5,gc2); pyxel.pset(gx2+3,gy2+3,7); pyxel.pset(gx2+5,gy2+3,7)
        pyxel.rect(gx2+2,gy2+7-walk,2,1,gc2); pyxel.rect(gx2+5,gy2+6+walk,1,1,gc2)
        pyxel.text(gx2-2,gy2+10,"G2",gc2)
        pyxel.line(14,70,SCREEN_W-14,70,5)
        # Left column
        lx=16
        pyxel.text(lx,75,"HOW TO PLAY",7)
        for txt,col,y in [("Arrows    Move & dig blocks",13,84),("P         Pickaxe upgrade",13,92),
                          ("O2        Oxygen tank refill",12,100),("v         Exit portal",10,108),("R         Return to menu",6,116)]:
            pyxel.text(lx,y,txt,col)
        pyxel.line(14,128,SCREEN_W//2-2,128,5)
        pyxel.text(lx,133,"ORES  (harder = deeper)",7)
        pyxel.rect(lx,142,6,6,4);  pyxel.text(lx+9,143,"Dirt     pickaxe 1",4)
        pyxel.rect(lx,152,6,6,13); pyxel.text(lx+9,153,"Stone    pickaxe 2",13)
        pyxel.rect(lx,162,6,6,6);  pyxel.text(lx+9,163,"Granite  pickaxe 3",6)
        pyxel.line(14,174,SCREEN_W//2-2,174,5)
        pyxel.text(lx,179,"RANKS  (fewer blocks = higher)",7)
        for i,(lo,hi,lbl,rc) in enumerate(RANKS[:6]):
            pyxel.text(lx,188+i*8,f"{lo}-{hi}  {lbl}",rc)
        # Right column
        rx=SCREEN_W//2+4
        pyxel.text(rx,75,"HAZARDS",8)
        pyxel.text(rx,84,"^  Spikes    instant death",8)
        pyxel.text(rx,92,"G  Gas leak  drains O2",3)
        pyxel.rect(rx,100,6,6,gc1); pyxel.text(rx+9,101,"G1: red, steady pace",gc1)
        pyxel.rect(rx,110,6,6,2);   pyxel.text(rx+9,111,"G2: purple, faster",2)
        pyxel.line(SCREEN_W//2+2,122,SCREEN_W-14,122,5)
        pyxel.text(rx,127,"3 LEVELS",7)
        pyxel.text(rx,136,"1 Surface    gentle start",7)
        pyxel.text(rx,144,"2 Midcore    faster O2",9)
        pyxel.text(rx,152,"3 Deepcore   brutal granite",8)
        pyxel.line(SCREEN_W//2+2,162,SCREEN_W-14,162,5)
        pyxel.text(rx,167,"GOAL",7)
        pyxel.text(rx,176,"Mine as FEW blocks as",6)
        pyxel.text(rx,184,"possible — low score wins!",6)
        pyxel.text(rx,192,"Fewer blocks = better rank",6)
        pyxel.line(SCREEN_W//2,70,SCREEN_W//2,228,5)
        pyxel.line(14,228,SCREEN_W-14,228,5)
        # Record + tip
        if self.best_blocks is not None:
            rl,rc=get_rank(self.best_blocks)
            pyxel.text(lx,233,f"Record: {self.best_blocks} blk  [{rl}]",rc)
        else:
            pyxel.text(lx,233,"No record yet — be the first!",6)
        tips=["Tip: lure guardians into dead ends!",
              "Tip: G2 is faster — dodge it first.",
              "Tip: grab O2 tanks before going deep.",
              "Tip: gas leaks drain O2 through walls."]
        pyxel.text(lx,212,tips[(f//180)%4],13 if (f//30)%2==0 else 7)
        if (f//25)%2==0:
            pyxel.rect(SCREEN_W//2-68,220,136,10,0); pyxel.rectb(SCREEN_W//2-68,220,136,10,10)
            pyxel.text(SCREEN_W//2-60,222,"Press  ENTER  or  SPACE  to start",10)

    # ── Draw ───────────────────────────────────────────────────────────────
    def draw(self):
        if self.game_state=="menu": self._menu(); return
        f=pyxel.frame_count
        pyxel.cls(LEVELS[self.level_idx]["bg"])
        # Tiles
        for y in range(H):
            for x in range(W):
                if not self.disc[y][x]: pyxel.rect(x*TILE,y*TILE,TILE,TILE,1); continue
                blk=self.world[y][x]
                if blk: draw_tile(x,y,blk)
        # Exit portal
        ex,ey=self.exit
        if self.disc[ey][ex]:
            if not self.picks:
                p=(f//8)%4; pc=[10,9,8,9]
                pyxel.rectb(ex*TILE-1,ey*TILE-1,TILE+2,TILE+2,pc[p])
                pyxel.rect(ex*TILE,ey*TILE,TILE,TILE,pc[(p+2)%4])
                pyxel.text(ex*TILE+1,ey*TILE+1,"v",7)
            else: pyxel.rectb(ex*TILE,ey*TILE,TILE,TILE,5)
        # O2 tanks
        for tx,ty in self.tanks:
            if self.disc[ty][tx]:
                pyxel.rect(tx*TILE+1,ty*TILE+1,TILE-2,TILE-2,12)
                pyxel.text(tx*TILE+1,ty*TILE+2,"O2",1)
        # Spikes
        for sx2,sy2 in self.spikes:
            if self.disc[sy2][sx2]:
                bx,by=sx2*TILE,sy2*TILE
                pyxel.pset(bx+2,by+6,8); pyxel.pset(bx+4,by+5,8); pyxel.pset(bx+6,by+6,8)
                pyxel.line(bx+2,by+7,bx+6,by+7,8)
        # Gas leaks
        for gx,gy in self.gas:
            if self.disc[gy][gx]:
                gc=[3,11,3][(f//6)%3]; bx,by=gx*TILE,gy*TILE
                pyxel.rectb(bx,by,TILE,TILE,gc); pyxel.text(bx+1,by+1,"G",gc)
            for r in range(1,3):
                for ddx in range(-r,r+1):
                    for ddy in range(-r,r+1):
                        nx2,ny2=gx+ddx,gy+ddy
                        if 0<=nx2r2:
                    bx,by=x*TILE,y*TILE
                    for oy in range(TILE):
                        for ox in range(TILE):
                            if (ox+oy)%2==0: pyxel.pset(bx+ox,by+oy,1)
        # Critical O2 red flash
        if self.o2.5 else(9 if ox>.25 else 8)
        if ox<.2 and (f//10)%2==0: pyxel.rect(SCREEN_W-bw-4,2,bw,7,8)
        else:                       pyxel.rect(SCREEN_W-bw-4,2,bw,7,1)
        pyxel.rect(SCREEN_W-bw-4,2,fil,7,oc); pyxel.text(SCREEN_W-bw-4,2,"O2",7)
        if self.mt>0:
            pyxel.rect(0,14,len(self.msg)*4+6,10,0); pyxel.text(4,16,self.msg,10)
        # Legend
        pyxel.rect(0,SCREEN_H-11,SCREEN_W,11,0)
        bt=f"Best:{self.best_blocks}blk" if self.best_blocks else "Best:--"
        pyxel.text(4,SCREEN_H-9,"Arrows:move  R=menu  ^=spike  G=gas",13)
        pyxel.text(SCREEN_W-len(bt)*4-4,SCREEN_H-9,bt,9)
        # Transition
        if self.trans:
            nxt=LEVELS[self.level_idx+1]["name"] if self.level_idx+1