"""
Models - Core data structures for MMA orchestration and project configuration.

This module defines the primary dataclasses used throughout the Manual Slop
application for representing tasks, tracks, and execution context.

Key Data Structures:
 - Ticket: Atomic unit of work with status, dependencies, and context requirements
 - Track: Collection of tickets with a shared goal
 - WorkerContext: Execution context for a Tier 3 worker
 - Metadata: Track metadata (id, name, status, timestamps)
 - TrackState: Serializable track state with discussion history
 - FileItem: File configuration with auto-aggregate and force-full flags

Status Machine (Ticket):
 todo -> in_progress -> completed
   |          |
   v          v
 blocked   blocked

Serialization:
 All dataclasses provide to_dict() and from_dict() class methods for TOML/JSON
 persistence via project_manager.py.

Thread Safety:
 These dataclasses are NOT thread-safe. Callers must synchronize mutations
 if sharing instances across threads (e.g., during ConductorEngine execution).

Configuration Integration:
 - load_config() / save_config() read/write the global config.toml
 - AGENT_TOOL_NAMES defines the canonical list of MCP tools available to agents

See Also:
 - docs/guide_mma.md for MMA orchestration documentation
 - src/dag_engine.py for TrackDAG and ExecutionEngine
 - src/multi_agent_conductor.py for ConductorEngine
 - src/project_manager.py for persistence layer
"""
from __future__ import annotations
import json
import os
import tomllib
import datetime
from dataclasses import dataclass, field
from typing import List, Optional, Dict, Any, Union
from pathlib import Path
from src.paths import get_config_path

CONFIG_PATH = get_config_path()

def _clean_nones(data: Any) -> Any:
 if isinstance(data, dict):
  return {k: _clean_nones(v) for k, v in data.items() if v is not None}
 elif isinstance(data, list):
  return [_clean_nones(v) for v in data if v is not None]
 return data

def load_config() -> dict[str, Any]:
 with open(CONFIG_PATH, "rb") as f:
  return tomllib.load(f)

def save_config(config: dict[str, Any]) -> None:
 import tomli_w
 import sys
 config = _clean_nones(config)
 sys.stderr.write(f"[DEBUG] Saving config. Theme: {config.get('theme')}\n")
 sys.stderr.flush()
 with open(CONFIG_PATH, "wb") as f:
  tomli_w.dump(config, f)

AGENT_TOOL_NAMES = [
 "run_powershell",
 "read_file",
 "list_directory",
 "search_files",
 "web_search",
 "fetch_url",
 "get_file_summary",
 "py_get_skeleton",
 "py_get_code_outline",
 "py_get_definition",
 "py_get_signature",
 "py_get_class_summary",
 "py_get_var_declaration",
 "py_get_docstring",
 "py_find_usages",
 "py_get_imports",
 "py_check_syntax",
 "py_get_hierarchy",
 "ts_c_get_skeleton",
 "ts_cpp_get_skeleton",
 "ts_c_get_code_outline",
 "ts_cpp_get_code_outline"
]

DEFAULT_TOOL_CATEGORIES: Dict[str, List[str]] = {
 "General": ["read_file", "list_directory", "search_files", "get_tree", "get_file_summary"],
 "Python": [
  "py_get_skeleton",
  "py_get_code_outline",
  "py_get_definition",
  "py_get_signature",
  "py_get_class_summary",
  "py_get_var_declaration",
  "py_get_docstring",
  "py_update_definition",
  "py_set_signature",
  "py_set_var_declaration"
 ],
 "Surgical": ["get_file_slice", "set_file_slice", "edit_file"],
 "Web": ["web_search", "fetch_url"],
 "Analysis": ["py_find_usages", "py_get_imports", "py_check_syntax", "py_get_hierarchy"],
 "C/C++": [
  "ts_c_get_skeleton",
  "ts_cpp_get_skeleton",
  "ts_c_get_code_outline",
  "ts_cpp_get_code_outline",
  "ts_c_get_definition",
  "ts_cpp_get_definition",
  "ts_c_get_signature",
  "ts_cpp_get_signature",
  "ts_c_update_definition",
  "ts_cpp_update_definition"
 ],
 "Runtime": ["run_powershell", "get_ui_performance"]
}

def parse_history_entries(history_strings: list[str], roles: list[str]) -> list[dict[str, Any]]:
 import re
 from src import thinking_parser
 entries = []
 for raw in history_strings:
  ts = ""
  rest = raw
  if rest.startswith("@"):
   nl = rest.find("\n")
   if nl != -1:
    ts = rest[1:nl]
    rest = rest[nl + 1:]
  known = roles or ["User", "AI", "Vendor API", "System"]
  role_pat = re.compile(r"^(" + "|".join(re.escape(r) for r in known) + r"):", re.IGNORECASE)
  match = role_pat.match(rest)
  role = match.group(1) if match else "User"
  if match:
   content = rest[match.end():].strip()
  else:
   content = rest
   
  entry_obj = {"role": role, "content": content, "collapsed": True, "ts": ts}
  if role == "AI" and ("<thinking>" in content or "<thought>" in content or "Thinking:" in content):
   segments, parsed_content = thinking_parser.parse_thinking_trace(content)
   if segments:
    entry_obj["content"] = parsed_content
    entry_obj["thinking_segments"] = [{"content": s.content, "marker": s.marker} for s in segments]
    
  entries.append(entry_obj)
 return entries

@dataclass
class ThinkingSegment:
 content: str
 marker: str # 'thinking', 'thought', or 'Thinking:'

 def to_dict(self) -> Dict[str, Any]:
  return {"content": self.content, "marker": self.marker}

 @classmethod
 def from_dict(cls, data: Dict[str, Any]) -> "ThinkingSegment":
  return cls(content=data["content"], marker=data["marker"])


@dataclass
class Ticket:
 id: str
 description: str
 target_symbols: List[str] = field(default_factory=list)
 context_requirements: List[str] = field(default_factory=list)
 depends_on: List[str] = field(default_factory=list)
 status: str = "todo"
 assigned_to: str = "unassigned"
 priority: str = "medium"
 target_file: Optional[str] = None
 blocked_reason: Optional[str] = None
 step_mode: bool = False
 retry_count: int = 0
 manual_block: bool = False
 model_override: Optional[str] = None
 persona_id: Optional[str] = None

 def mark_blocked(self, reason: str) -> None:
  self.status = "blocked"
  self.blocked_reason = reason

 def mark_manual_block(self, reason: str) -> None:
  self.status = "blocked"
  self.blocked_reason = f"[MANUAL] {reason}"
  self.manual_block = True

 def clear_manual_block(self) -> None:
  if self.manual_block:
   self.status = "todo"
   self.blocked_reason = None
   self.manual_block = False

 def mark_complete(self) -> None:
  self.status = "completed"

 def get(self, key: str, default: Any = None) -> Any:
  return getattr(self, key, default)

 def to_dict(self) -> Dict[str, Any]:
  return {
   "id": self.id,
   "description": self.description,
   "status": self.status,
   "assigned_to": self.assigned_to,
   "priority": self.priority,
   "target_file": self.target_file,
   "target_symbols": self.target_symbols,
   "context_requirements": self.context_requirements,
   "depends_on": self.depends_on,
   "blocked_reason": self.blocked_reason,
   "step_mode": self.step_mode,
   "retry_count": self.retry_count,
   "manual_block": self.manual_block,
   "model_override": self.model_override,
   "persona_id": self.persona_id,
  }

 @classmethod
 def from_dict(cls, data: Dict[str, Any]) -> "Ticket":
  return cls(
   id=data["id"],
   description=data.get("description", ""),
   status=data.get("status", "todo"),
   assigned_to=data.get("assigned_to", "unassigned"),
   priority=data.get("priority", "medium"),
   target_file=data.get("target_file"),
   target_symbols=data.get("target_symbols", []),
   context_requirements=data.get("context_requirements", []),
   depends_on=data.get("depends_on", []),
   blocked_reason=data.get("blocked_reason"),
   step_mode=data.get("step_mode", False),
   retry_count=data.get("retry_count", 0),
   manual_block=data.get("manual_block", False),
   model_override=data.get("model_override"),
   persona_id=data.get("persona_id"),
  )


@dataclass
class Track:
 id: str
 description: str
 tickets: List[Ticket] = field(default_factory=list)

 def get_executable_tickets(self) -> List[Ticket]:
  from src.dag_engine import TrackDAG
  dag = TrackDAG(self.tickets)
  return dag.get_ready_tasks()

 def to_dict(self) -> Dict[str, Any]:
  return {
   "id": self.id,
   "description": self.description,
   "tickets": [t.to_dict() for t in self.tickets],
  }

 @classmethod
 def from_dict(cls, data: Dict[str, Any]) -> "Track":
  return cls(
   id=data["id"],
   description=data.get("description", ""),
   tickets=[Ticket.from_dict(t) for t in data.get("tickets", [])],
  )


@dataclass
class WorkerContext:
 ticket_id: str
 model_name: str
 messages: List[Dict[str, Any]] = field(default_factory=list)
 tool_preset: Optional[str] = None
 persona_id: Optional[str] = None


@dataclass
class Metadata:
 id: str
 name: str
 status: Optional[str] = None
 created_at: Optional[datetime.datetime] = None
 updated_at: Optional[datetime.datetime] = None

 def to_dict(self) -> Dict[str, Any]:
  return {
   "id": self.id,
   "name": self.name,
   "status": self.status,
   "created_at": self.created_at.isoformat() if self.created_at else None,
   "updated_at": self.updated_at.isoformat() if self.updated_at else None,
  }

 @classmethod
 def from_dict(cls, data: Dict[str, Any]) -> "Metadata":
  created = data.get("created_at")
  updated = data.get("updated_at")
  if isinstance(created, str):
   try:
    created = datetime.datetime.fromisoformat(created)
   except ValueError:
    created = None
  if isinstance(updated, str):
   try:
    updated = datetime.datetime.fromisoformat(updated)
   except ValueError:
    updated = None
  return cls(
   id=data["id"],
   name=data.get("name", ""),
   status=data.get("status"),
   created_at=created,
   updated_at=updated,
  )


@dataclass
class TrackState:
 metadata: Metadata
 discussion: List[str] = field(default_factory=list)
 tasks: List[Ticket] = field(default_factory=list)

 def to_dict(self) -> Dict[str, Any]:
  serialized_discussion = []
  for item in self.discussion:
   if isinstance(item, dict):
    new_item = dict(item)
    if "ts" in new_item and isinstance(new_item["ts"], datetime.datetime):
     new_item["ts"] = new_item["ts"].isoformat()
    serialized_discussion.append(new_item)
   else:
    serialized_discussion.append(item)
  return {
   "metadata": self.metadata.to_dict(),
   "discussion": serialized_discussion,
   "tasks": [t.to_dict() for t in self.tasks],
  }

 @classmethod
 def from_dict(cls, data: Dict[str, Any]) -> "TrackState":
  discussion = data.get("discussion", [])
  parsed_discussion = []
  for item in discussion:
   if isinstance(item, dict):
    new_item = dict(item)
    ts = new_item.get("ts")
    if isinstance(ts, str):
     try:
      new_item["ts"] = datetime.datetime.fromisoformat(ts)
     except ValueError:
      pass
    parsed_discussion.append(new_item)
   else:
    parsed_discussion.append(item)
  return cls(
   metadata=Metadata.from_dict(data["metadata"]),
   discussion=parsed_discussion,
   tasks=[Ticket.from_dict(t) for t in data.get("tasks", [])],
  )

@dataclass
class FileItem:
 path: str
 auto_aggregate: bool = True
 force_full: bool = False
 injected_at: Optional[float] = None

 def to_dict(self) -> Dict[str, Any]:
  return {
   "path": self.path,
   "auto_aggregate": self.auto_aggregate,
   "force_full": self.force_full,
   "injected_at": self.injected_at,
  }

 @classmethod
 def from_dict(cls, data: Dict[str, Any]) -> "FileItem":
  return cls(
   path=data["path"],
   auto_aggregate=data.get("auto_aggregate", True),
   force_full=data.get("force_full", False),
   injected_at=data.get("injected_at"),
  )

@dataclass
class Preset:
 name: str
 system_prompt: str

 def to_dict(self) -> Dict[str, Any]:
  return {
   "system_prompt": self.system_prompt,
  }

 @classmethod
 def from_dict(cls, name: str, data: Dict[str, Any]) -> "Preset":
  return cls(
   name=name,
   system_prompt=data.get("system_prompt", ""),
  )

@dataclass
class Tool:
 name: str
 approval: str = 'auto'
 weight: int = 3
 parameter_bias: Dict[str, str] = field(default_factory=dict)

 def to_dict(self) -> Dict[str, Any]:
  return {
   "name": self.name,
   "approval": self.approval,
   "weight": self.weight,
   "parameter_bias": self.parameter_bias,
  }

 @classmethod
 def from_dict(cls, data: Dict[str, Any]) -> "Tool":
  return cls(
   name=data["name"],
   approval=data.get("approval", "auto"),
   weight=data.get("weight", 3),
   parameter_bias=data.get("parameter_bias", {}),
  )

@dataclass
class ToolPreset:
 name: str
 categories: Dict[str, List[Union[Tool, Any]]]

 def to_dict(self) -> Dict[str, Any]:
  serialized_categories = {}
  for cat, tools in self.categories.items():
   serialized_categories[cat] = [t.to_dict() if isinstance(t, Tool) else t for t in tools]
  return {
   "categories": serialized_categories,
  }

 @classmethod
 def from_dict(cls, name: str, data: Dict[str, Any]) -> "ToolPreset":
  raw_categories = data.get("categories", {})
  parsed_categories = {}
  for cat, tools in raw_categories.items():
   parsed_categories[cat] = [Tool.from_dict(t) if isinstance(t, dict) else t for t in tools]
  return cls(
   name=name,
   categories=parsed_categories,
  )

@dataclass
class BiasProfile:
 name: str
 tool_weights: Dict[str, int] = field(default_factory=dict)
 category_multipliers: Dict[str, float] = field(default_factory=dict)

 def to_dict(self) -> Dict[str, Any]:
  return {
   "name": self.name,
   "tool_weights": self.tool_weights,
   "category_multipliers": self.category_multipliers,
  }

 @classmethod
 def from_dict(cls, data: Dict[str, Any]) -> "BiasProfile":
  return cls(
   name=data["name"],
   tool_weights=data.get("tool_weights", {}),
   category_multipliers=data.get("category_multipliers", {}),
  )

@dataclass
class Persona:
 name: str
 preferred_models: List[Dict[str, Any]] = field(default_factory=list)
 system_prompt: str = ''
 tool_preset: Optional[str] = None
 bias_profile: Optional[str] = None
 context_preset: Optional[str] = None
 aggregation_strategy: Optional[str] = None

 @property
 def provider(self) -> Optional[str]:
  if not self.preferred_models: return None
  return self.preferred_models[0].get("provider")

 @property
 def model(self) -> Optional[str]:
  if not self.preferred_models: return None
  return self.preferred_models[0].get("model")

 @property
 def temperature(self) -> Optional[float]:
  if not self.preferred_models: return None
  return self.preferred_models[0].get("temperature")

 @property
 def top_p(self) -> Optional[float]:
  if not self.preferred_models: return None
  return self.preferred_models[0].get("top_p")

 @property
 def max_output_tokens(self) -> Optional[int]:
  if not self.preferred_models: return None
  return self.preferred_models[0].get("max_output_tokens")

 def to_dict(self) -> Dict[str, Any]:
  res = {
   "system_prompt": self.system_prompt,
  }
  if self.preferred_models:
   processed = []
   for m in self.preferred_models:
    if isinstance(m, str):
     processed.append({"model": m})
    else:
     processed.append(m)
   res["preferred_models"] = processed
  if self.tool_preset is not None:
   res["tool_preset"] = self.tool_preset
  if self.bias_profile is not None:
   res["bias_profile"] = self.bias_profile
  if self.context_preset is not None:
   res["context_preset"] = self.context_preset
  if self.aggregation_strategy is not None:
   res["aggregation_strategy"] = self.aggregation_strategy
  return res

 @classmethod
 def from_dict(cls, name: str, data: Dict[str, Any]) -> "Persona":
  raw_models = data.get("preferred_models", [])
  parsed_models = []
  for m in raw_models:
   if isinstance(m, str):
    parsed_models.append({"model": m})
   else:
    parsed_models.append(m)

  # Migration logic: merge legacy fields if they exist
  legacy = {}
  for k in ["provider", "model", "temperature", "top_p", "max_output_tokens"]:
   if data.get(k) is not None:
    legacy[k] = data[k]

  if legacy:
   if not parsed_models:
    parsed_models.append(legacy)
   else:
    # Merge into first item if it's missing these specific legacy fields
    for k, v in legacy.items():
     if k not in parsed_models[0] or parsed_models[0][k] is None:
      parsed_models[0][k] = v

  return cls(
   name=name,
   preferred_models=parsed_models,
   system_prompt=data.get("system_prompt", ""),
   tool_preset=data.get("tool_preset"),
   bias_profile=data.get("bias_profile"),
   context_preset=data.get("context_preset"),
   aggregation_strategy=data.get("aggregation_strategy"),
  )
@dataclass
class MCPServerConfig:
 name: str
 command: Optional[str] = None
 args: List[str] = field(default_factory=list)
 url: Optional[str] = None
 auto_start: bool = False

 def to_dict(self) -> Dict[str, Any]:
  res = {'auto_start': self.auto_start}
  if self.command: res['command'] = self.command
  if self.args: res['args'] = self.args
  if self.url: res['url'] = self.url
  return res

 @classmethod
 def from_dict(cls, name: str, data: Dict[str, Any]) -> 'MCPServerConfig':
  return cls(
   name=name,
   command=data.get('command'),
   args=data.get('args', []),
   url=data.get('url'),
   auto_start=data.get('auto_start', False),
  )

@dataclass
class MCPConfiguration:
 mcpServers: Dict[str, MCPServerConfig] = field(default_factory=dict)

 def to_dict(self) -> Dict[str, Any]:
  return {
   'mcpServers': {name: cfg.to_dict() for name, cfg in self.mcpServers.items()}
  }

 @classmethod
 def from_dict(cls, data: Dict[str, Any]) -> 'MCPConfiguration':
  raw_servers = data.get('mcpServers', {})
  parsed_servers = {
   name: MCPServerConfig.from_dict(name, cfg)
   for name, cfg in raw_servers.items()
  }
  return cls(mcpServers=parsed_servers)

@dataclass
class VectorStoreConfig:
 provider: str # 'chroma', 'qdrant', 'mock', 'mcp'
 url: Optional[str] = None
 api_key: Optional[str] = None
 collection_name: str = 'manual_slop'
 mcp_server: Optional[str] = None
 mcp_tool: Optional[str] = None

 def to_dict(self) -> Dict[str, Any]:
  return {
   "provider": self.provider,
   "url": self.url,
   "api_key": self.api_key,
   "collection_name": self.collection_name,
   "mcp_server": self.mcp_server,
   "mcp_tool": self.mcp_tool,
  }

 @classmethod
 def from_dict(cls, data: Dict[str, Any]) -> "VectorStoreConfig":
  return cls(
   provider=data["provider"],
   url=data.get("url"),
   api_key=data.get("api_key"),
   collection_name=data.get("collection_name", "manual_slop"),
   mcp_server=data.get("mcp_server"),
   mcp_tool=data.get("mcp_tool"),
  )

@dataclass
class RAGConfig:
 enabled: bool = False
 vector_store: VectorStoreConfig = field(default_factory=lambda: VectorStoreConfig(provider='mock'))
 embedding_provider: str = 'gemini'
 chunk_size: int = 1000
 chunk_overlap: int = 200

 def to_dict(self) -> Dict[str, Any]:
  return {
   "enabled": self.enabled,
   "vector_store": self.vector_store.to_dict(),
   "embedding_provider": self.embedding_provider,
   "chunk_size": self.chunk_size,
   "chunk_overlap": self.chunk_overlap,
  }

 @classmethod
 def from_dict(cls, data: Dict[str, Any]) -> "RAGConfig":
  return cls(
   enabled=data.get("enabled", False),
   vector_store=VectorStoreConfig.from_dict(data.get("vector_store", {"provider": "mock"})),
   embedding_provider=data.get("embedding_provider", "gemini"),
   chunk_size=data.get("chunk_size", 1000),
   chunk_overlap=data.get("chunk_overlap", 200),
  )

@dataclass
class WorkspaceProfile:
 name: str
 ini_content: str
 show_windows: Dict[str, bool]
 panel_states: Dict[str, Any]

 def to_dict(self) -> Dict[str, Any]:
  return {
   "ini_content": self.ini_content,
   "show_windows": self.show_windows,
   "panel_states": self.panel_states,
  }

 @classmethod
 def from_dict(cls, name: str, data: Dict[str, Any]) -> "WorkspaceProfile":
  return cls(
   name=name,
   ini_content=data.get("ini_content", ""),
   show_windows=data.get("show_windows", {}),
   panel_states=data.get("panel_states", {}),
  )

def load_mcp_config(path: str) -> MCPConfiguration:
 if not os.path.exists(path):
  return MCPConfiguration()
 with open(path, 'r', encoding='utf-8') as f:
  try:
   data = json.load(f)
   return MCPConfiguration.from_dict(data)
  except Exception:
   return MCPConfiguration()

