import tkinter as tk
import math, random, time, threading, socket

# ---------------- Config ----------------
WINDOW_W, WINDOW_H = 1280, 720
MAP_W, MAP_H       = 120, 60
BASE_TILE          = 16
VIEW_W, VIEW_H     = 60, 40
FPS                = 15
ZOOM_STEP          = 1.1
NET_PORT           = 50007

# ---------------- Tank Models & Stats ----------------
TANK_MODELS = {
    "USA": {
        "M4 Sherman":   ["   ____   ","  /____\\  "," | M4  |  "," |_____|  "],
        "M48 Patton":   ["   _____  ","  [|===|] ","  | M48 | ","  [|___|] "],
        "M1A1 Abrams":  ["   /--\\   ","  |====|  ","  |AbrSR| ","  \\____/  "],
    },
    "USSR": {
        "T-34":         ["   ____   ","  /____\\  "," | T-34|  "," |_____|  "],
        "T-72":         ["   <==>   ","  /====\\  "," | T-72|  ","  \\____/  "],
        "T-90":         ["   [==]   ","  |====|  "," | T-90|  ","  \\____/  "],
    },
    "China": {
        "T-26":         ["   .--.   ","  [|::|]  ","  | T26|  ","  |____|  "],
        "Type 59":      ["   /::\\   ","  |====|  ","  |59::|  ","   \\__/   "],
        "99A":          ["   <++>   ","  /++++\\  "," |  99A|  ","  \\____/  "],
    }
}
TANK_STATS = {
    "USA":{
      "M4 Sherman":(100,2,5),
      "M48 Patton":(120,1.5,5),
      "M1A1 Abrams":(90,2.5,4),
    },
    "USSR":{
      "T-34":(95,2.5,5),
      "T-72":(110,2.0,5),
      "T-90":(100,2.2,4),
    },
    "China":{
      "T-26":(80,2.8,6),
      "Type 59":(105,2.0,5),
      "99A":(95,2.5,4),
    }
}
INF_MODEL = [" I "]  # infantry

# ---------------- Explosion Styles ----------------
EXPLOSION_STYLES = {
  "USA":   (["   *   ","  ***  "," ***** ","*******"," ***** ","  ***  ","   *   "],"yellow"),
  "USSR":  (["   +   ","  +++  "," +++++ ","+++++++"," +++++ ","  +++  ","   +   "],"red"),
  "China": (["   %   ","  %%%  "," %%%%% ","%%%%%%%"," %%%%% ","  %%%  ","   %   "],"cyan"),
}

# ---------------- Utils ----------------
def dist(x1,y1,x2,y2): return math.hypot(x2-x1,y2-y1)

# ---------------- Classes ----------------
class Explosion:
    def __init__(self,x,y,country):
        frm, col = EXPLOSION_STYLES[country]
        self.frames, self.color = frm, col
        self.x, self.y = x,y
        self.i=self.t=0
    def update(self):
        self.t+=1
        if self.t>=1: self.t=0; self.i+=1
    def finished(self): return self.i>=len(self.frames)
    def frame(self):   return self.frames[self.i] if self.i<len(self.frames) else ""

class Shell:
    def __init__(self,x,y,ang,owner):
        self.x,self.y,self.ang,self.owner = x,y,ang,owner
        self.speed, self.damage = 3, 40
    def move(self):
        r=math.radians(self.ang)
        self.x += math.cos(r)*self.speed
        self.y += math.sin(r)*self.speed

class Tank:
    def __init__(self,country,name,x,y,is_player=False):
        self.country, self.name = country,name
        self.model = TANK_MODELS[country][name]
        self.max_hp,self.speed,self.max_ammo = TANK_STATS[country][name]
        self.hp,self.ammo = self.max_hp, self.max_ammo
        self.x,self.y,en = x,y,is_player
        self.turret,self.is_player,self.kills=0,is_player,0
        self.last_shot,self.fire_delay=0,0.5
    def alive(self): return self.hp>0

class Infantry:
    def __init__(self,uid,x,y):
        self.id,self.model = uid,INF_MODEL
        self.hp,self.speed=50,1.5
        self.x,self.y= x,y
        self.last_shot,self.fire_delay=0,1.5
    def alive(self): return self.hp>0

# ---------------- Networking ----------------
class Net(threading.Thread):
    def __init__(self,g,host,ip):
        super().__init__(daemon=True)
        self.g,self.host,self.ip = g,host,ip; self.sock=None; self.running=True
    def run(self):
        if self.host: self.server()
        else:         self.client()
    def server(self):
        s=socket.socket(); s.bind(("",NET_PORT)); s.listen(1)
        self.sock, _ = s.accept()
        while self.running:
            try:
                d=self.sock.recv(1024).decode()
                if d: self.g.recv_net(d)
                self.sock.sendall(self.g.send_net().encode())
            except: break
        s.close()
    def client(self):
        self.sock=socket.socket()
        try: self.sock.connect((self.ip,NET_PORT))
        except: return
        while self.running:
            try:
                d=self.sock.recv(1024).decode()
                if d: self.g.recv_net(d)
                self.sock.sendall(self.g.send_net().encode())
            except: break
        self.sock.close()

# ---------------- Main Menu ----------------
class MainMenu(tk.Frame):
    def __init__(self,root,cb):
        super().__init__(root); self.pack(fill="both",expand=True)
        self.cb,self.mode = cb, tk.StringVar("offline")
        self.country, self.tank = tk.StringVar("USA"), tk.StringVar("M4 Sherman")
        tk.Label(self,text="ASCII Tank Game",font=("Courier",24)).pack(pady=5)
        mf=tk.LabelFrame(self,text="Mode"); mf.pack(fill="x")
        for m in ["offline","local","host","join"]:
            tk.Radiobutton(mf,text=m.capitalize(),var=self.mode,value=m).pack(anchor="w")
        cf=tk.LabelFrame(self,text="Country"); cf.pack(fill="x")
        for c in TANK_MODELS: tk.Radiobutton(cf,text=c,var=self.country,value=c).pack(anchor="w")
        tf=tk.LabelFrame(self,text="Tank"); tf.pack(fill="x")
        self.tf,self.tf_widgets=tf,[]
        self.update_tanks()
        self.country.trace("w",lambda *a:self.update_tanks())
        tk.Button(self,text="Start",command=self.start).pack(pady=5)
    def update_tanks(self):
        for w in self.tf_widgets: w.destroy()
        self.tf_widgets.clear()
        for t in TANK_MODELS[self.country.get()]:
            w=tk.Radiobutton(self.tf,text=t,var=self.tank,value=t)
            w.pack(anchor="w"); self.tf_widgets.append(w)
    def start(self):
        ip="" 
        if self.mode.get()=="join":
            ip=tk.simpledialog.askstring("Host IP","Enter host IP") or ""
        self.cb(self.mode.get(),self.country.get(),
                self.tank.get(), self.mode.get()=="host", ip)

# ---------------- The Game ----------------
class GameApp:
    def __init__(self,root,mode,country,tank,host,ip):
        self.root, self.mode = root, mode
        self.canvas=tk.Canvas(root,width=WINDOW_W,height=WINDOW_H,bg="black")
        self.canvas.pack(); self.hud=tk.Label(root,font=("Courier",14),bg="gray20",fg="white")
        self.hud.pack(fill="x")
        # map
        self.map=[[ "." if random.random()>0.12 else random.choice("MBKo") for _ in range(MAP_W)] for _ in range(MAP_H)]
        # player
        self.p1=Tank(country,tank,MAP_W//2,MAP_H//2,True)
        # collections
        self.enemies={}  # enemy tanks & inf
        self.allies={}   # AI friends
        # spawn 5 enemies & 3 infantry
        for i in range(5):
            x,y = random.uniform(0,MAP_W), random.uniform(0,MAP_H)
            while dist(x,y,self.p1.x,self.p1.y)<20: x,y = random.uniform(0,MAP_W), random.uniform(0,MAP_H)
            name=random.choice(list(TANK_MODELS["USSR"].keys()))
            self.enemies[f"E{i}"]=Tank("USSR",name,x,y,False)
        for i in range(3):
            x,y = random.uniform(0,MAP_W), random.uniform(0,MAP_H)
            while dist(x,y,self.p1.x,self.p1.y)<20: x,y = random.uniform(0,MAP_W), random.uniform(0,MAP_H)
            self.enemies[f"I{i}"]=Infantry(f"I{i}",x,y)
        # spawn 2 AI allies near player
        for i in range(2):
            x,y = self.p1.x+random.uniform(-5,5), self.p1.y+random.uniform(-5,5)
            name=random.choice(list(TANK_MODELS[country].keys()))
            self.allies[f"A{i}"]=Tank(country,name,x,y,False)
        # other state
        self.shells=[]; self.exps=[]
        self.zoom=1; self.camx=self.camy=0; self.mx=self.my=0
        # network?
        if mode in ("host","join"):
            self.net=Net(self,host,ip); self.wait=True; self.show_loading(host)
        # bindings & loop
        root.bind("<Key>",self.on_key)
        self.canvas.bind("<Motion>",self.on_mouse)
        self.canvas.bind("<Button-1>",self.on_click)
        self.last=time.time(); self.loop()

    def show_loading(self,host):
        txt="Waiting for player..." if host else "Connecting..."
        self.canvas.create_text(WINDOW_W/2,WINDOW_H/2,text=txt,fill="white",font=("Courier",24))
        self.canvas.update()
        while getattr(self,"wait",False): time.sleep(0.1)
        self.canvas.delete("all")

    def on_key(self,e):
        k=e.keysym.lower()
        if k in("w","a","s","d"):
            m={"w":(0,-1),"s":(0,1),"a":(-1,0),"d":(1,0)}[k]
            self.p1.x+=m[0]*self.p1.speed; self.p1.y+=m[1]*self.p1.speed
        if self.mode=="local" and k in("up","down","left","right"):
            m={"up":(0,-1),"down":(0,1),"left":(-1,0),"right":(1,0)}[k]
            ally=list(self.allies.values())[0]  # assign first ally as second local player
            ally.x+=m[0]*ally.speed; ally.y+=m[1]*ally.speed
        if k=="space": self.fire(self.p1)
        if e.char=="+": self.zoom*=ZOOM_STEP
        if e.char=="-": self.zoom/=ZOOM_STEP
        self.clamp(); self.update_cam()

    def on_mouse(self,e):
        self.mx,self.my=e.x,e.y
        cx,cy=WINDOW_W/2,WINDOW_H/2
        self.p1.turret=math.degrees(math.atan2(e.y-cy,e.x-cx))

    def on_click(self,e): self.fire(self.p1)

    def clamp(self):
        self.p1.x= max(0,min(MAP_W-1,self.p1.x))
        self.p1.y= max(0,min(MAP_H-1,self.p1.y))

    def update_cam(self):
        self.camx=self.p1.x-VIEW_W/2; self.camy=self.p1.y-VIEW_H/2
        self.camx=max(0,min(MAP_W-VIEW_W,self.camx))
        self.camy=max(0,min(MAP_H-VIEW_H,self.camy))

    def fire(self,t):
        now=time.time()
        if hasattr(t,"ammo") and now-t.last_shot>t.fire_delay and t.ammo>0:
            self.shells.append(Shell(t.x,t.y,t.turret,t))
            t.ammo-=1; t.last_shot=now
            if self.mode in ("host","join"):
                self.net.sock.sendall(self.send_net().encode())

    def recv_net(self,data):
        self.wait=False
        pid,x,y,ang,hp,ammo,kills = data.strip().split(";")
        x,y,ang,hp=float(x),float(y),float(ang),float(hp)
        ammo,kills=int(ammo),int(kills)
        if pid not in self.allies:  # treat remote as ally
            self.allies[pid] = Tank("USSR","T-34",x,y,False)
        p=self.allies[pid]
        p.x,p.y,p.turret,p.hp,p.ammo,p.kills = x,y,ang,hp,ammo,kills

    def send_net(self):
        p=self.p1
        return f"PLAYER;{p.x};{p.y};{p.turret};{p.hp};{p.ammo};{p.kills}\n"

    def update(self,dt):
        # update shells
        for s in self.shells[:]:
            s.move()
            if not(0<=s.x<MAP_W and 0<=s.y<MAP_H):
                self.shells.remove(s); continue
            # collisions vs all tanks & inf
            for col in list(self.enemies.values())+list(self.allies.values()):
                if hasattr(col,"hp") and col.alive() and dist(s.x,s.y,col.x,col.y)<1:
                    col.hp -= s.damage
                    self.shells.remove(s)
                    self.exps.append(Explosion(s.x,s.y,col.country if hasattr(col,"country") else "USA"))
                    if col.hp<=0 and s.owner.is_player: s.owner.kills+=1
                    break
        # explosions
        for e in self.exps[:]:
            e.update()
            if e.finished(): self.exps.remove(e)
        # enemy AI
        for en in self.enemies.values():
            if en.alive():
                dx,dy = self.p1.x-en.x, self.p1.y-en.y
                en.turret = math.degrees(math.atan2(dy,dx))
                if dist(self.p1.x,self.p1.y,en.x,en.y)<20 and time.time()-en.last_shot>en.fire_delay:
                    self.fire(en); en.last_shot=time.time()
        # ally AI (non-player)
        for al in self.allies.values():
            if not al.is_player and al.alive():
                # follow player loosely
                dx,dy = self.p1.x-al.x, self.p1.y-al.y
                if dist(self.p1.x,self.p1.y,al.x,al.y)>5:
                    ang=math.atan2(dy,dx)
                    al.x+=math.cos(ang)*al.speed*0.5
                    al.y+=math.sin(ang)*al.speed*0.5
                # shoot nearest enemy
                target = min(self.enemies.values(), key=lambda e: dist(al.x,al.y,e.x,e.y))
                dx,dy=target.x-al.x, target.y-al.y
                al.turret=math.degrees(math.atan2(dy,dx))
                if dist(al.x,al.y,target.x,target.y)<20 and time.time()-al.last_shot>al.fire_delay:
                    self.fire(al); al.last_shot=time.time()

    def draw(self):
        self.canvas.delete("all")
        z=self.zoom; ts=BASE_TILE*z; tilt=4*z
        # map
        for j in range(VIEW_H):
            for i in range(VIEW_W):
                mx,my = int(self.camx+i), int(self.camy+j)
                if 0<=mx<MAP_W and 0<=my<MAP_H:
                    c=self.map[my][mx]
                    col={".":"#8c8","M":"#aaa","B":"#fc0","K":"#555","o":"#642"}[c]
                    x=i*ts+j*tilt*0.2
                    y=j*ts-j*0.5
                    self.canvas.create_text(x,y,text=c,fill=col,font=("Courier",int(12*z)))
        # draw tanks & inf
        for unit in [self.p1]+list(self.enemies.values())+list(self.allies.values()):
            if hasattr(unit,"hp") and not unit.alive(): continue
            d = dist(self.p1.x,self.p1.y,unit.x,unit.y)
            sx=(unit.x-self.camx)*ts + (unit.y-self.camy)*tilt*0.2
            sy=(unit.y-self.camy)*ts
            # far
            if d>10:
                sym = "T" if hasattr(unit,"model") else "i"
                col = "cyan" if unit.is_player else ("green" if unit in self.allies.values() else "magenta")
                self.canvas.create_text(sx,sy,text=sym,fill=col,
                    font=("Courier",max(8,int(12*z*0.5))))
            else:
                if hasattr(unit,"model"):
                    for idx,line in enumerate(unit.model):
                        self.canvas.create_text(sx,sy+idx*12*z,text=line,
                            fill="cyan" if unit.is_player else ("green" if unit in self.allies.values() else "red"),
                            font=("Courier",int(12*z)))
                    # turret
                    cx,cy = sx+len(unit.model[0])*6*z/2, sy+len(unit.model)*12*z/2
                    r=20*z; rad=math.radians(unit.turret)
                    self.canvas.create_line(cx,cy,cx+math.cos(rad)*r, cy+math.sin(rad)*r,
                                            fill="cyan" if unit.is_player else "green", width=2)
                else:
                    self.canvas.create_text(sx,sy,text="i",fill="green",font=("Courier",int(12*z)))
            # health bar
            if hasattr(unit,"hp"):
                w,h = 40*z,5
                r=unit.hp/unit.max_hp
                self.canvas.create_rectangle(sx,sy-8*z,sx+w,sy-8*z+h,fill="black")
                self.canvas.create_rectangle(sx,sy-8*z,sx+w*r,sy-8*z+h,
                    fill="green" if r>0.4 else "yellow" if r>0.2 else "red")
        # shells
        for s in self.shells:
            x=(s.x-self.camx)*ts; y=(s.y-self.camy)*ts
            self.canvas.create_text(x,y,text="*",fill="white",font=("Courier",int(12*z)))
        # explosions
        for e in self.exps:
            f=e.frame(); x=(e.x-self.camx)*ts; y=(e.y-self.camy)*ts
            self.canvas.create_text(x,y,text=f,fill=e.color,font=("Courier",int(12*z)))
        # crosshair
        self.canvas.create_oval(self.mx-5,self.my-5,self.mx+5,self.my+5,outline="lime")
        # HUD
        p=self.p1
        self.canvas.create_text(10,WINDOW_H-20,anchor="nw",
            text=f"HP:{p.hp} Ammo:{p.ammo} Kills:{p.kills} Zoom:{self.zoom:.2f}",
            fill="white",font=("Courier",14))

    def loop(self):
        now=time.time(); dt=now-self.last; self.last=now
        if not getattr(self,"wait",False):
            self.update(dt); self.update_cam(); self.draw()
        self.canvas.after(int(1000/FPS),self.loop)

# ---------------- Startup ----------------
def start(mode,country,tank,host,ip):
    menu.destroy()
    g = GameApp(root,mode,country,tank,host,ip)
    g.last = time.time()

root=tk.Tk()
menu = MainMenu(root, start)
root.mainloop()