"""
MundiX CLI - Project / Client Workspace Manager
Manages per-client/target workspaces with full context persistence.

Each project creates a structured workspace:
    ~/mundix-projects/<project_name>/
        ├── project.json          # Project metadata
        ├── context.json          # Rolling context (last N command outputs)
        ├── findings.json         # Discovered findings (ports, services, vulns)
        ├── timeline.json         # Full timeline of actions
        ├── loot/                 # Captured files, hashes, credentials
        ├── scans/                # Raw scan outputs
        ├── reports/              # Generated reports
        └── notes.md              # User notes

The context system feeds command outputs back to the AI so it "remembers"
what happened in the engagement.
"""

import json
import os
import re
import time
from datetime import datetime
from pathlib import Path
from typing import Optional, List, Dict

PROJECTS_BASE = Path.home() / "mundix-projects"
MUNDIX_DIR = Path.home() / ".mundix"
ACTIVE_PROJECT_FILE = MUNDIX_DIR / "active_project.json"

# Maximum context entries to feed to the AI (rolling window)
MAX_CONTEXT_ENTRIES = 20
# Maximum chars per command output stored in context
MAX_OUTPUT_CHARS = 4000


class Project:
    """Manages a single pentest project/client workspace."""

    def __init__(self, name: str = None):
        self.name = name
        self.project_dir = None
        self.metadata = {}
        self.context = {"entries": [], "summary": ""}
        self.findings = {"hosts": {}, "services": [], "vulnerabilities": [], "credentials": [], "notes": []}
        self.timeline = {"events": []}

        if name:
            self._load_or_create(name)
        else:
            self._try_load_active()

    # ─── Project Lifecycle ──────────────────────────────────────────

    def _try_load_active(self):
        """Try to load the last active project."""
        try:
            if ACTIVE_PROJECT_FILE.exists():
                data = json.loads(ACTIVE_PROJECT_FILE.read_text())
                active_name = data.get("active_project")
                if active_name:
                    project_dir = PROJECTS_BASE / active_name
                    if project_dir.exists():
                        self.name = active_name
                        self._load_or_create(active_name)
        except Exception:
            pass

    def _load_or_create(self, name: str):
        """Load existing project or create new one."""
        self.name = self._sanitize_name(name)
        self.project_dir = PROJECTS_BASE / self.name

        if self.project_dir.exists():
            self._load_project()
        else:
            self._create_project()

        self._set_active()

    def _sanitize_name(self, name: str) -> str:
        """Sanitize project name for filesystem."""
        sanitized = re.sub(r'[^\w\-.]', '_', name.lower().strip())
        return sanitized[:64]

    def _create_project(self):
        """Create a new project workspace."""
        self.project_dir.mkdir(parents=True, exist_ok=True)
        (self.project_dir / "loot").mkdir(exist_ok=True)
        (self.project_dir / "scans").mkdir(exist_ok=True)
        (self.project_dir / "reports").mkdir(exist_ok=True)

        self.metadata = {
            "name": self.name,
            "created_at": datetime.now().isoformat(),
            "updated_at": datetime.now().isoformat(),
            "target": "",
            "scope": [],
            "type": "",  # web, internal, app
            "status": "active",
            "description": "",
        }

        self.context = {"entries": [], "summary": ""}
        self.findings = {
            "hosts": {},
            "services": [],
            "vulnerabilities": [],
            "credentials": [],
            "notes": [],
        }
        self.timeline = {"events": []}

        # Create initial notes file
        notes_path = self.project_dir / "notes.md"
        notes_path.write_text(
            f"# {self.name}\n\n"
            f"**Created:** {datetime.now().strftime('%Y-%m-%d %H:%M')}\n\n"
            f"## Notes\n\n"
            f"_Add your notes here..._\n"
        )

        self._save_all()

    def _load_project(self):
        """Load an existing project from disk."""
        self.metadata = self._load_json("project.json", self.metadata)
        self.context = self._load_json("context.json", self.context)
        self.findings = self._load_json("findings.json", self.findings)
        self.timeline = self._load_json("timeline.json", self.timeline)

    def _set_active(self):
        """Set this project as the active one."""
        MUNDIX_DIR.mkdir(parents=True, exist_ok=True)
        ACTIVE_PROJECT_FILE.write_text(json.dumps({
            "active_project": self.name,
            "set_at": datetime.now().isoformat(),
        }))

    def _save_all(self):
        """Save all project data."""
        self.metadata["updated_at"] = datetime.now().isoformat()
        self._save_json("project.json", self.metadata)
        self._save_json("context.json", self.context)
        self._save_json("findings.json", self.findings)
        self._save_json("timeline.json", self.timeline)

    # ─── JSON Helpers ───────────────────────────────────────────────

    def _load_json(self, filename: str, default: dict) -> dict:
        """Load a JSON file from the project directory."""
        path = self.project_dir / filename
        try:
            if path.exists():
                return json.loads(path.read_text())
        except (json.JSONDecodeError, IOError):
            pass
        return default

    def _save_json(self, filename: str, data: dict):
        """Save data to a JSON file in the project directory."""
        path = self.project_dir / filename
        try:
            path.write_text(json.dumps(data, indent=2, ensure_ascii=False, default=str))
        except IOError as e:
            print(f"[!] Failed to save {filename}: {e}")

    # ─── Context Management (THE KEY FEATURE) ───────────────────────

    def add_context(self, entry_type: str, content: str, command: str = None,
                    exit_code: int = None, metadata: dict = None):
        """
        Add an entry to the rolling context window.
        This is what feeds the AI with "memory" of what happened.

        Args:
            entry_type: "command_output", "ai_response", "user_query", "finding"
            content: The actual content (truncated if too long)
            command: The command that was executed (if applicable)
            exit_code: Command exit code (if applicable)
            metadata: Additional metadata
        """
        entry = {
            "id": len(self.context["entries"]) + 1,
            "timestamp": datetime.now().isoformat(),
            "type": entry_type,
            "content": content[:MAX_OUTPUT_CHARS],
            "command": command,
            "exit_code": exit_code,
            "metadata": metadata or {},
        }

        self.context["entries"].append(entry)

        # Rolling window: keep only the last N entries
        if len(self.context["entries"]) > MAX_CONTEXT_ENTRIES * 2:
            # Keep a summary of older entries + recent ones
            old_entries = self.context["entries"][:-MAX_CONTEXT_ENTRIES]
            self.context["entries"] = self.context["entries"][-MAX_CONTEXT_ENTRIES:]

            # Build summary of what was trimmed
            old_summary_parts = []
            for e in old_entries:
                if e["type"] == "command_output" and e.get("command"):
                    status = "OK" if e.get("exit_code", 0) == 0 else f"FAIL({e.get('exit_code')})"
                    old_summary_parts.append(f"- [{status}] {e['command']}")
                elif e["type"] == "finding":
                    old_summary_parts.append(f"- FINDING: {e['content'][:100]}")

            if old_summary_parts:
                existing_summary = self.context.get("summary", "")
                new_summary = "\n".join(old_summary_parts[-20:])
                self.context["summary"] = (existing_summary + "\n" + new_summary)[-2000:]

        self._save_json("context.json", self.context)

        # Also add to timeline
        self._add_timeline_event(entry_type, command or content[:100], metadata)

    def get_context_for_ai(self) -> str:
        """
        Build a context string to inject into the AI prompt.
        This is what makes the AI "remember" previous results.
        """
        if not self.name:
            return ""

        parts = []

        # Project info
        parts.append(f"=== ACTIVE PROJECT: {self.name} ===")
        if self.metadata.get("target"):
            parts.append(f"Target: {self.metadata['target']}")
        if self.metadata.get("scope"):
            parts.append(f"Scope: {', '.join(self.metadata['scope'])}")
        if self.metadata.get("type"):
            parts.append(f"Type: {self.metadata['type']}")

        # Summary of older context (if any)
        if self.context.get("summary"):
            parts.append(f"\n--- Previous Actions Summary ---\n{self.context['summary']}")

        # Recent context entries (the rolling window)
        recent = self.context["entries"][-MAX_CONTEXT_ENTRIES:]
        if recent:
            parts.append("\n--- Recent Session Context ---")
            for entry in recent:
                ts = entry.get("timestamp", "")[:19]
                etype = entry.get("type", "unknown")

                if etype == "command_output":
                    cmd = entry.get("command", "?")
                    exit_code = entry.get("exit_code", "?")
                    output = entry.get("content", "")
                    # Truncate very long outputs for context
                    if len(output) > 1500:
                        output = output[:750] + "\n[...truncated...]\n" + output[-750:]
                    parts.append(f"\n[CMD] $ {cmd} (exit: {exit_code})\n{output}")

                elif etype == "user_query":
                    parts.append(f"\n[USER] {entry.get('content', '')[:200]}")

                elif etype == "ai_response":
                    parts.append(f"\n[AI] {entry.get('content', '')[:300]}")

                elif etype == "finding":
                    parts.append(f"\n[FINDING] {entry.get('content', '')}")

        # Known findings summary
        findings_summary = self._get_findings_summary()
        if findings_summary:
            parts.append(f"\n--- Known Findings ---\n{findings_summary}")

        return "\n".join(parts)

    def _get_findings_summary(self) -> str:
        """Build a summary of discovered findings."""
        parts = []

        # Hosts
        hosts = self.findings.get("hosts", {})
        if hosts:
            parts.append("Hosts discovered:")
            for ip, info in list(hosts.items())[:10]:
                ports = info.get("ports", [])
                port_str = ", ".join([f"{p['port']}/{p.get('service', '?')}" for p in ports[:8]])
                parts.append(f"  {ip}: {port_str}")

        # Vulnerabilities
        vulns = self.findings.get("vulnerabilities", [])
        if vulns:
            parts.append(f"Vulnerabilities ({len(vulns)}):")
            for v in vulns[:5]:
                parts.append(f"  - {v.get('title', 'Unknown')}: {v.get('severity', '?')}")

        # Credentials
        creds = self.findings.get("credentials", [])
        if creds:
            parts.append(f"Credentials found: {len(creds)}")

        return "\n".join(parts) if parts else ""

    # ─── Findings Management ────────────────────────────────────────

    def add_host(self, ip: str, hostname: str = None, os_info: str = None):
        """Add or update a discovered host."""
        if ip not in self.findings["hosts"]:
            self.findings["hosts"][ip] = {
                "hostname": hostname,
                "os": os_info,
                "ports": [],
                "discovered_at": datetime.now().isoformat(),
            }
        else:
            if hostname:
                self.findings["hosts"][ip]["hostname"] = hostname
            if os_info:
                self.findings["hosts"][ip]["os"] = os_info
        self._save_json("findings.json", self.findings)

    def add_port(self, ip: str, port: int, protocol: str = "tcp",
                 service: str = None, version: str = None, state: str = "open"):
        """Add a discovered port to a host."""
        if ip not in self.findings["hosts"]:
            self.add_host(ip)

        port_entry = {
            "port": port,
            "protocol": protocol,
            "service": service,
            "version": version,
            "state": state,
        }

        # Avoid duplicates
        existing_ports = self.findings["hosts"][ip]["ports"]
        for i, p in enumerate(existing_ports):
            if p["port"] == port and p["protocol"] == protocol:
                existing_ports[i] = port_entry
                self._save_json("findings.json", self.findings)
                return

        existing_ports.append(port_entry)
        self._save_json("findings.json", self.findings)

    def add_vulnerability(self, title: str, severity: str = "medium",
                          host: str = None, port: int = None, description: str = None,
                          cve: str = None, evidence: str = None):
        """Add a discovered vulnerability."""
        vuln = {
            "id": len(self.findings["vulnerabilities"]) + 1,
            "title": title,
            "severity": severity,
            "host": host,
            "port": port,
            "description": description,
            "cve": cve,
            "evidence": evidence,
            "discovered_at": datetime.now().isoformat(),
        }
        self.findings["vulnerabilities"].append(vuln)
        self._save_json("findings.json", self.findings)

    def add_credential(self, username: str, password: str = None,
                       hash_value: str = None, source: str = None):
        """Add a discovered credential."""
        cred = {
            "username": username,
            "password": password,
            "hash": hash_value,
            "source": source,
            "discovered_at": datetime.now().isoformat(),
        }
        self.findings["credentials"].append(cred)
        self._save_json("findings.json", self.findings)

    def add_note(self, note: str):
        """Add a note to findings."""
        self.findings["notes"].append({
            "content": note,
            "timestamp": datetime.now().isoformat(),
        })
        self._save_json("findings.json", self.findings)

    # ─── Auto-Parse Scan Results ────────────────────────────────────

    def parse_nmap_output(self, output: str, target_ip: str = None):
        """Auto-parse nmap output to extract hosts, ports, services."""
        # Extract host IPs
        ip_pattern = r'Nmap scan report for (?:(\S+) \()?(\d+\.\d+\.\d+\.\d+)\)?'
        for match in re.finditer(ip_pattern, output):
            hostname = match.group(1)
            ip = match.group(2)
            self.add_host(ip, hostname=hostname)

            # Extract ports for this host
            # Look for port lines after this host
            host_section_start = match.end()
            next_host = re.search(r'Nmap scan report for', output[host_section_start:])
            host_section = output[host_section_start:host_section_start + next_host.start() if next_host else len(output)]

            port_pattern = r'(\d+)/(tcp|udp)\s+(\w+)\s+(\S+)(?:\s+(.+))?'
            for port_match in re.finditer(port_pattern, host_section):
                port_num = int(port_match.group(1))
                protocol = port_match.group(2)
                state = port_match.group(3)
                service = port_match.group(4)
                version = port_match.group(5) if port_match.group(5) else None

                if state == "open":
                    self.add_port(ip, port_num, protocol, service, version, state)

            # Extract OS info
            os_pattern = r'OS details?:\s*(.+)'
            os_match = re.search(os_pattern, host_section)
            if os_match:
                self.add_host(ip, os_info=os_match.group(1).strip())

        # If no host found but target_ip provided
        if target_ip and target_ip not in self.findings["hosts"]:
            self.add_host(target_ip)
            port_pattern = r'(\d+)/(tcp|udp)\s+(\w+)\s+(\S+)(?:\s+(.+))?'
            for port_match in re.finditer(port_pattern, output):
                port_num = int(port_match.group(1))
                protocol = port_match.group(2)
                state = port_match.group(3)
                service = port_match.group(4)
                version = port_match.group(5) if port_match.group(5) else None
                if state == "open":
                    self.add_port(target_ip, port_num, protocol, service, version, state)

    def save_scan_output(self, command: str, output: str, filename: str = None):
        """Save raw scan output to the scans directory."""
        if not self.project_dir:
            return None

        scans_dir = self.project_dir / "scans"
        scans_dir.mkdir(exist_ok=True)

        if not filename:
            timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
            tool = command.split()[0].replace("/", "_") if command else "unknown"
            filename = f"{tool}_{timestamp}.txt"

        filepath = scans_dir / filename
        filepath.write_text(f"# Command: {command}\n# Timestamp: {datetime.now().isoformat()}\n\n{output}")
        return str(filepath)

    # ─── Timeline ───────────────────────────────────────────────────

    def _add_timeline_event(self, event_type: str, description: str, metadata: dict = None):
        """Add an event to the project timeline."""
        event = {
            "timestamp": datetime.now().isoformat(),
            "type": event_type,
            "description": description[:200],
            "metadata": metadata or {},
        }
        self.timeline["events"].append(event)

        # Keep timeline manageable
        if len(self.timeline["events"]) > 500:
            self.timeline["events"] = self.timeline["events"][-300:]

        self._save_json("timeline.json", self.timeline)

    # ─── Report Generation ──────────────────────────────────────────

    def generate_report(self) -> str:
        """Generate a Markdown pentest report from project data."""
        if not self.project_dir:
            return "No active project."

        report_parts = []

        # Header
        report_parts.append(f"# Pentest Report: {self.name}")
        report_parts.append(f"\n**Generated:** {datetime.now().strftime('%Y-%m-%d %H:%M')}")
        report_parts.append(f"**Target:** {self.metadata.get('target', 'N/A')}")
        report_parts.append(f"**Type:** {self.metadata.get('type', 'N/A')}")
        report_parts.append(f"**Scope:** {', '.join(self.metadata.get('scope', ['N/A']))}")
        report_parts.append(f"**Status:** {self.metadata.get('status', 'active')}")

        # Executive Summary
        report_parts.append("\n## Executive Summary\n")
        hosts_count = len(self.findings.get("hosts", {}))
        vulns = self.findings.get("vulnerabilities", [])
        creds = self.findings.get("credentials", [])
        total_ports = sum(len(h.get("ports", [])) for h in self.findings.get("hosts", {}).values())

        report_parts.append(
            f"This assessment identified **{hosts_count} host(s)** with "
            f"**{total_ports} open port(s)**, **{len(vulns)} vulnerability(ies)**, "
            f"and **{len(creds)} credential(s)** discovered."
        )

        # Severity breakdown
        severity_counts = {}
        for v in vulns:
            sev = v.get("severity", "info").lower()
            severity_counts[sev] = severity_counts.get(sev, 0) + 1
        if severity_counts:
            report_parts.append("\n| Severity | Count |")
            report_parts.append("| --- | --- |")
            for sev in ["critical", "high", "medium", "low", "info"]:
                if sev in severity_counts:
                    report_parts.append(f"| {sev.upper()} | {severity_counts[sev]} |")

        # Hosts & Services
        report_parts.append("\n## Hosts & Services\n")
        for ip, info in self.findings.get("hosts", {}).items():
            hostname = info.get("hostname", "")
            os_info = info.get("os", "Unknown")
            report_parts.append(f"### {ip}" + (f" ({hostname})" if hostname else ""))
            report_parts.append(f"**OS:** {os_info}\n")

            ports = info.get("ports", [])
            if ports:
                report_parts.append("| Port | Protocol | State | Service | Version |")
                report_parts.append("| --- | --- | --- | --- | --- |")
                for p in sorted(ports, key=lambda x: x["port"]):
                    report_parts.append(
                        f"| {p['port']} | {p.get('protocol', 'tcp')} | "
                        f"{p.get('state', 'open')} | {p.get('service', '?')} | "
                        f"{p.get('version', '-')} |"
                    )
            report_parts.append("")

        # Vulnerabilities
        if vulns:
            report_parts.append("\n## Vulnerabilities\n")
            for v in sorted(vulns, key=lambda x: ["critical", "high", "medium", "low", "info"].index(x.get("severity", "info").lower()) if x.get("severity", "info").lower() in ["critical", "high", "medium", "low", "info"] else 4):
                report_parts.append(f"### [{v.get('severity', '?').upper()}] {v.get('title', 'Unknown')}")
                if v.get("host"):
                    report_parts.append(f"**Host:** {v['host']}" + (f":{v['port']}" if v.get('port') else ""))
                if v.get("cve"):
                    report_parts.append(f"**CVE:** {v['cve']}")
                if v.get("description"):
                    report_parts.append(f"\n{v['description']}")
                if v.get("evidence"):
                    report_parts.append(f"\n**Evidence:**\n```\n{v['evidence'][:500]}\n```")
                report_parts.append("")

        # Credentials
        if creds:
            report_parts.append("\n## Credentials Found\n")
            report_parts.append("| Username | Password/Hash | Source |")
            report_parts.append("| --- | --- | --- |")
            for c in creds:
                pw = c.get("password") or c.get("hash", "N/A")
                report_parts.append(f"| {c.get('username', '?')} | `{pw}` | {c.get('source', '?')} |")

        # Timeline
        report_parts.append("\n## Activity Timeline\n")
        events = self.timeline.get("events", [])[-30:]
        if events:
            report_parts.append("| Time | Type | Description |")
            report_parts.append("| --- | --- | --- |")
            for e in events:
                ts = e.get("timestamp", "")[:19]
                report_parts.append(f"| {ts} | {e.get('type', '?')} | {e.get('description', '')[:80]} |")

        # Notes
        notes = self.findings.get("notes", [])
        if notes:
            report_parts.append("\n## Notes\n")
            for n in notes:
                report_parts.append(f"- [{n.get('timestamp', '')[:19]}] {n.get('content', '')}")

        report_parts.append(f"\n---\n*Report generated by MundiX CLI | {datetime.now().isoformat()}*\n")

        report_content = "\n".join(report_parts)

        # Save report
        reports_dir = self.project_dir / "reports"
        reports_dir.mkdir(exist_ok=True)
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        report_path = reports_dir / f"report_{timestamp}.md"
        report_path.write_text(report_content)

        return str(report_path)

    # ─── Project Info ───────────────────────────────────────────────

    def get_info(self) -> dict:
        """Get project summary info."""
        if not self.name:
            return {"active": False}

        return {
            "active": True,
            "name": self.name,
            "target": self.metadata.get("target", ""),
            "type": self.metadata.get("type", ""),
            "status": self.metadata.get("status", "active"),
            "hosts": len(self.findings.get("hosts", {})),
            "ports": sum(len(h.get("ports", [])) for h in self.findings.get("hosts", {}).values()),
            "vulns": len(self.findings.get("vulnerabilities", [])),
            "creds": len(self.findings.get("credentials", [])),
            "context_entries": len(self.context.get("entries", [])),
            "timeline_events": len(self.timeline.get("events", [])),
            "path": str(self.project_dir),
            "created": self.metadata.get("created_at", ""),
        }

    def set_target(self, target: str):
        """Set the project target."""
        self.metadata["target"] = target
        if target not in self.metadata.get("scope", []):
            self.metadata.setdefault("scope", []).append(target)
        self._save_json("project.json", self.metadata)

    def set_type(self, project_type: str):
        """Set the project type (web, internal, app)."""
        self.metadata["type"] = project_type
        self._save_json("project.json", self.metadata)

    def close(self):
        """Mark project as completed."""
        self.metadata["status"] = "completed"
        self.metadata["completed_at"] = datetime.now().isoformat()
        self._save_all()

    @staticmethod
    def list_projects() -> List[dict]:
        """List all existing projects."""
        projects = []
        if PROJECTS_BASE.exists():
            for d in sorted(PROJECTS_BASE.iterdir()):
                if d.is_dir() and (d / "project.json").exists():
                    try:
                        meta = json.loads((d / "project.json").read_text())
                        projects.append({
                            "name": meta.get("name", d.name),
                            "target": meta.get("target", ""),
                            "type": meta.get("type", ""),
                            "status": meta.get("status", "active"),
                            "created": meta.get("created_at", ""),
                        })
                    except Exception:
                        projects.append({"name": d.name, "status": "unknown"})
        return projects

    @staticmethod
    def delete_project(name: str) -> bool:
        """Delete a project workspace."""
        import shutil
        project_dir = PROJECTS_BASE / name
        if project_dir.exists():
            shutil.rmtree(project_dir)
            return True
        return False
