import io
import re
import time
import json
import base64
import random
import string
import logging
import subprocess
from typing import Optional

import requests
from PIL import Image  # pillow for tiny PNG generation

from .Challenge import Challenge


class Blogpost(Challenge):
    """
    SLA checker for the provided Flask challenge app.

    Assumptions / Defaults (override as needed):
      - Service listens on self.port (inherited)
      - Container name is 'chal_app' (override via self.container_name)
      - Flag lives in '/app/flag.txt' inside container
      - Login flow at /login, register at /register
      - Create post at /create, view post at /post/<id>, uploads at /uploads/<file>
    """
    flag_location = 'flags/blogpost.txt'           # Host copy (used by your orchestrator)
    history_location = 'history/blogpost.txt'
    container_flag_path = '/flag.txt'
    container_name = 'blogpost_container'                  # <-- set to your actual container name

    # Heuristics to recognize ExifTool output
    _exif_markers = (
        'ExifTool Version Number',
        'File Name',
        'File Size',
        'MIME Type',
        'File Type',
    )

    def _make_logger(self):
        if not hasattr(self, 'logger') or self.logger is None:
            self.logger = logging.getLogger(self.__class__.__name__)
            if not self.logger.handlers:
                h = logging.StreamHandler()
                h.setFormatter(logging.Formatter('[%(levelname)s] %(message)s'))
                self.logger.addHandler(h)
                self.logger.setLevel(logging.INFO)

    # --- Flag distribution hook (optional, mirrors your example) ---
    def distribute(self, flag: str) -> bool:
        """
        Writes/records the current flag on the host. Your infra may separately
        mount/copy it into the container; this class *also* verifies existence
        inside the container during .check().
        """
        self._make_logger()
        try:
            with open(self.flag_location, 'w') as f:
                f.write(flag)
            with open(self.history_location, 'a') as f:
                f.write(flag + '\n')
            self.logger.info(f"Flag '{flag}' written to {self.flag_location}")
            return True
        except Exception as e:
            self.logger.error(f"Failed writing host flag: {e}")
            return False

    # --- Helpers ---
    def _gen_username(self) -> str:
        return "user_" + ''.join(random.choices(string.ascii_lowercase + string.digits, k=8))

    def _gen_password(self) -> str:
        return "Pw!" + ''.join(random.choices(string.ascii_letters + string.digits, k=10))

    def _tiny_png_bytes(self) -> bytes:
        """
        Generate a minimal valid PNG in-memory to trigger exiftool.
        """
        img = Image.new("RGB", (2, 2), (123, 200, 50))
        buf = io.BytesIO()
        img.save(buf, format="PNG")
        return buf.getvalue()

    def _docker_exec(self, args: list[str], timeout: int = 10) -> subprocess.CompletedProcess:
        """
        Run `docker exec` into the challenge container.
        """
        return subprocess.run(
            ["docker", "exec", self.container_name] + args,
            capture_output=True,
            text=True,
            timeout=timeout
        )

    # --- SLA core ---
    def check(self) -> bool:
        self._make_logger()
        base_url = f"http://localhost:{self.port}"
        s = requests.Session()

        # 0) Liveness: login page should be reachable (no auth needed)
        login_url = base_url + "/login"
        self.logger.info(f"[1/7] Checking liveness at {login_url} ...")
        try:
            r = s.get(login_url, timeout=10)
            assert r.status_code == 200, f"Login page HTTP {r.status_code}"
            self.logger.info("    ✓ Login page reachable")
        except Exception as e:
            self.logger.error(f"    ✗ Liveness check failed: {e}")
            return False

        # 1) Register a fresh user
        self.logger.info("[2/7] Registering a fresh user ...")
        username = self._gen_username()
        password = self._gen_password()
        try:
            r = s.post(
                base_url + "/register",
                data={"username": username, "password": password},
                allow_redirects=False,
                timeout=10,
            )
            # Flask typically redirects to /login on success (302)
            assert r.status_code in (200, 302, 303), f"Register HTTP {r.status_code}"
            self.logger.info(f"    ✓ Registered as {username}")
        except Exception as e:
            self.logger.error(f"    ✗ Registration failed: {e}")
            return False

        # 2) Log in
        self.logger.info("[3/7] Logging in ...")
        try:
            r = s.post(
                base_url + "/login",
                data={"username": username, "password": password},
                allow_redirects=True,
                timeout=10,
            )
            # Successful login should redirect or render index; ensure not bounced back to /login with "Invalid credentials"
            assert r.status_code == 200, f"Login HTTP {r.status_code}"
            assert "/login" not in r.url or "Invalid credentials" not in r.text
            self.logger.info("    ✓ Logged in")
        except Exception as e:
            self.logger.error(f"    ✗ Login failed: {e}")
            return False

        # 3) Create post with image upload (PNG) to trigger exiftool + DB write
        self.logger.info("[4/7] Creating a post with image upload ...")
        title = "SLA Post " + ''.join(random.choices(string.ascii_letters, k=6))
        content = "hello_from_sla_checker"
        png_bytes = self._tiny_png_bytes()
        files = {
            "image": ("probe.png", png_bytes, "image/png")
        }
        data = {"title": title, "content": content}
        try:
            r = s.post(base_url + "/create", files=files, data=data, allow_redirects=True, timeout=20)
            assert r.status_code == 200, f"Create returned HTTP {r.status_code}"
            self.logger.info("    ✓ Post created")
        except Exception as e:
            self.logger.error(f"    ✗ Create post failed: {e}")
            return False

        # 4) Discover the newly created post id by scraping links like /post/<id> from index
        self.logger.info("[5/7] Resolving new post id from index ...")
        try:
            r = s.get(base_url + "/", timeout=10)
            assert r.status_code == 200, f"Index HTTP {r.status_code}"
            # Find the most recent /post/<id> that likely corresponds to our post title
            # First try to find blocks containing our title then pull an id:
            post_id: Optional[int] = None
            # Prefer links sitting near our title
            title_blocks = [m.start() for m in re.finditer(re.escape(title), r.text)]
            if title_blocks:
                # Search backward/forward around the title occurrences for a /post/<num> link
                for pos in title_blocks:
                    window = r.text[max(0, pos-1000):pos+1000]
                    m = re.search(r'/post/(\d+)', window)
                    if m:
                        post_id = int(m.group(1))
                        break
            # Fallback: pick the first /post/<id> in the page
            if post_id is None:
                m = re.search(r'/post/(\d+)', r.text)
                assert m, "No /post/<id> link found on index"
                post_id = int(m.group(1))

            self.logger.info(f"    ✓ Resolved post id = {post_id}")
        except Exception as e:
            self.logger.error(f"    ✗ Could not resolve post id: {e}")
            return False

        # 5) Visit the post page, confirm it renders and contains our title/content.
        self.logger.info("[6/7] Visiting the post page ...")
        try:
            vp = s.get(base_url + f"/post/{post_id}", timeout=10)
            assert vp.status_code == 200, f"Post page HTTP {vp.status_code}"
            assert title in vp.text, "Post title not present on view page"
            assert content in vp.text, "Post content not present on view page"
            self.logger.info("    ✓ Post page renders with our content")
        except Exception as e:
            self.logger.error(f"    ✗ View post failed: {e}")
            return False

        # 6) Verify the uploaded image is accessible and check metadata
        self.logger.info("[7/7] Verifying uploaded image and metadata ...")
        try:
            # Find something like /uploads/<sha256>.png (or .jpg/.jpeg/.bmp)
            m = re.search(r'/uploads/([A-Za-z0-9_.-]+\.(?:png|jpg|jpeg|bmp))', vp.text, flags=re.IGNORECASE)
            assert m, "No uploaded image link found on post page"
            image_name = m.group(1)
            self.logger.info(f"    → Found image: {image_name}")
            
            # Verify the image itself is accessible
            img_url = base_url + f"/uploads/{image_name}"
            img_r = s.get(img_url, timeout=10)
            assert img_r.status_code == 200, f"Image file HTTP {img_r.status_code}"
            assert len(img_r.content) > 0, "Image file is empty"
            self.logger.info("    ✓ Uploaded image accessible")
            
            # Check if metadata file exists
            meta_url = base_url + f"/uploads/{image_name}.meta"
            self.logger.info(f"    → Trying metadata at: {meta_url}")
            mr = s.get(meta_url, timeout=10)
            if mr.status_code == 200:
                meta_text = mr.text.strip()
                if any(tag in meta_text for tag in self._exif_markers):
                    self.logger.info("    ✓ Exif metadata present and readable")
                else:
                    self.logger.warning(f"    ⚠ Metadata file exists but doesn't look like ExifTool output")
            else:
                # Try without .meta extension, maybe it's embedded or stored differently
                self.logger.warning(f"    ⚠ Metadata file returned HTTP {mr.status_code}")
                # Non-fatal - as long as upload/display works
        except Exception as e:
            self.logger.error(f"    ✗ Upload verification failed: {e}")
            return False

        # 7) Flag existence in container (do not fail SLA if only host copy exists but container is missing—treat as warning or policy-driven)
        try:
            proc = self._docker_exec(["/bin/sh", "-lc", f"test -f {self.container_flag_path} && cat {self.container_flag_path} || echo __MISSING__"])
            out = (proc.stdout or "").strip()
            if "__MISSING__" in out or proc.returncode not in (0,):
                self.logger.warning("⚠ Flag file missing inside container")
            else:
                self.logger.info("    ✓ Container flag present")
                # Optional: compare with host flag if present
                try:
                    with open(self.flag_location, "r") as f:
                        host_flag = f.read().strip()
                    if host_flag and host_flag == out:
                        self.logger.info("    ✓ Host and container flags match")
                    else:
                        self.logger.warning("⚠ Host/container flag mismatch (may be expected if rotated separately)")
                except FileNotFoundError:
                    self.logger.warning("⚠ Host flag not found; skipping comparison")
        except Exception as e:
            # Non-fatal: you can tune this to fail the round if flag is mandatory.
            self.logger.warning(f"Flag existence check encountered an issue: {e}")

        self.logger.info("SLA check passed ✅")
        return True