// Package workspaceops defines the deliberately closed host workspace contract.
package workspaceops

import (
	"bytes"
	"context"
	"crypto/sha256"
	"encoding/base64"
	"encoding/hex"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"reflect"
	"regexp"
	"strings"
	"unicode/utf8"

	"github.com/aritmolab/thothii/tools/thothctl/internal/compose"
	"github.com/aritmolab/thothii/tools/thothctl/internal/config"
	"github.com/aritmolab/thothii/tools/thothctl/internal/safeio"
	"gopkg.in/yaml.v3"
)

const CodeRegistryBootstrapRecoveryConflict = "registry_bootstrap_recovery_conflict"
const (
	maxSQLFile         = 1 << 20
	maxSQLTotal        = 16 << 20
	maxAssumptions     = 256
	maxAssumptionBytes = 256
	maxResult          = 1 << 20
	// Requests carry up to 16 MiB of raw SQL or annotation bytes encoded as base64.
	// This is deliberately independent from maxResult, which bounds child stdout.
	maxRequest     = 24 << 20
	maxCandidate   = 700 << 10
	maxAnnotations = 16 << 20
)

var workspaceIDPattern = regexp.MustCompile(`^[a-z][a-z0-9-]{2,62}$`)
var runIDPattern = regexp.MustCompile(`^[0-9a-f]{32}$`)
var revisionPattern = regexp.MustCompile(`^[0-9a-f]{40}$`)
var digestPattern = regexp.MustCompile(`^sha256:[0-9a-f]{64}$`)

// Command is intentionally closed: callers cannot provide a child argv or bootstrap run ID.
type Command interface{ workspaceCommand() }
type InspectCommand struct {
	WorkspaceID string
	JSON        bool
}
type DwhRequest struct {
	WorkspaceID, Resume string
	JSON                bool
}
type SuggestFksRequest struct {
	WorkspaceID string
	FromSQL     []string
	Assume      []string
	Output      string
	JSON        bool
}
type CheckSchemaRequest struct {
	WorkspaceID, Resume, Annotations, ReviewedCandidates string
	JSON                                                 bool
}
type IndexSchemaRequest struct {
	WorkspaceID string
	JSON        bool
}
type EvidenceRequest struct {
	WorkspaceID string
	DryRun      bool
	Resume      string
	JSON        bool
}
type RunRequest struct {
	WorkspaceID, Resume string
	JSON                bool
}

func (InspectCommand) workspaceCommand()     {}
func (DwhRequest) workspaceCommand()         {}
func (SuggestFksRequest) workspaceCommand()  {}
func (CheckSchemaRequest) workspaceCommand() {}
func (IndexSchemaRequest) workspaceCommand() {}
func (EvidenceRequest) workspaceCommand()    {}
func (RunRequest) workspaceCommand()         {}

type Result struct {
	SchemaVersion      int                `json:"schemaVersion"`
	Status             string             `json:"status"`
	Code               string             `json:"code"`
	WorkspaceID        string             `json:"workspaceId"`
	WorkspaceRevision  string             `json:"workspaceRevision"`
	DescriptorBlob     string             `json:"descriptorBlob"`
	Operation          string             `json:"operation"`
	RunID              string             `json:"runId,omitempty"`
	ChildRuns          map[string]string  `json:"childRuns,omitempty"`
	CompletedStages    []string           `json:"completedStages"`
	Counts             map[string]int     `json:"counts,omitempty"`
	ArtifactIdentities []ArtifactIdentity `json:"artifactIdentities,omitempty"`
	Warnings           []string           `json:"warnings,omitempty"`
}
type ArtifactIdentity struct {
	Kind   string `json:"kind"`
	Digest string `json:"digest"`
}
type hostExport struct {
	MediaType     string `json:"mediaType"`
	SHA256        string `json:"sha256"`
	ContentBase64 string `json:"contentBase64"`
}

type sqlInput struct {
	Basename      string `json:"basename"`
	ContentBase64 string `json:"contentBase64"`
	SHA256        string `json:"sha256"`
}
type annotationInput struct {
	Basename      string `json:"basename"`
	ContentBase64 string `json:"contentBase64"`
	SHA256        string `json:"sha256"`
}
type inputEnvelope struct {
	SchemaVersion      int              `json:"schemaVersion"`
	Operation          string           `json:"operation"`
	WorkspaceID        string           `json:"workspaceId"`
	Resume             string           `json:"resume,omitempty"`
	SQL                []sqlInput       `json:"sql,omitempty"`
	Assume             []string         `json:"assume,omitempty"`
	Annotations        *annotationInput `json:"annotations,omitempty"`
	ReviewedCandidates string           `json:"reviewedCandidates,omitempty"`
	DryRun             bool             `json:"dryRun,omitempty"`
}

func invalid(msg string) (Command, error) { return nil, errors.New(msg) }
func requireWorkspace(v string) error {
	if !workspaceIDPattern.MatchString(v) {
		return errors.New("--workspace must be a valid workspace ID")
	}
	return nil
}
func requireRun(v string) error {
	if !runIDPattern.MatchString(v) {
		return errors.New("--resume must be a lowercase 32-hex run ID")
	}
	return nil
}
func one(args []string, i *int, flag string) (string, error) {
	if *i+1 >= len(args) || strings.HasPrefix(args[*i+1], "--") {
		return "", fmt.Errorf("%s requires a value", flag)
	}
	*i++
	return args[*i], nil
}

func ParseWorkspaceCommand(args []string) (Command, error) {
	if len(args) < 2 || args[0] != "workspace" {
		return invalid("workspace command is required")
	}
	// Parse the hierarchy before options. This prevents aliases such as workspace schema dwh.
	path := strings.Join(args[:min(3, len(args))], " ")
	var action string
	var rest []string
	switch {
	case args[1] == "inspect":
		action = "inspect"
		rest = args[2:]
	case len(args) >= 3 && args[1] == "preprocess" && (args[2] == "dwh" || args[2] == "evidence" || args[2] == "run"):
		action = "preprocess " + args[2]
		rest = args[3:]
	case len(args) >= 3 && args[1] == "schema" && (args[2] == "suggest-fks" || args[2] == "check"):
		action = "schema " + args[2]
		rest = args[3:]
	case args[1] == "index-schema":
		action = "index-schema"
		rest = args[2:]
	default:
		return invalid("unknown workspace command: " + path)
	}
	var ws, resume, annotations, reviewed, output string
	var jsonOut, dry bool
	var sql, assume []string
	seen := map[string]bool{}
	allowed := func(flag string) bool {
		switch action {
		case "inspect":
			return flag == "--workspace" || flag == "--json"
		case "preprocess dwh", "preprocess run":
			return flag == "--workspace" || flag == "--resume" || flag == "--json"
		case "preprocess evidence":
			return flag == "--workspace" || flag == "--resume" || flag == "--dry-run" || flag == "--json"
		case "schema suggest-fks":
			return flag == "--workspace" || flag == "--from-sql" || flag == "--assume" || flag == "--output" || flag == "--json"
		case "schema check":
			return flag == "--workspace" || flag == "--resume" || flag == "--annotations" || flag == "--reviewed-candidates" || flag == "--json"
		case "index-schema":
			return flag == "--workspace" || flag == "--json"
		}
		return false
	}
	for i := 0; i < len(rest); i++ {
		f := rest[i]
		if !allowed(f) {
			return invalid("unknown or invalid workspace option")
		}
		if f == "--json" || f == "--dry-run" {
			if seen[f] {
				return invalid("duplicate " + f)
			}
			seen[f] = true
			if f == "--json" {
				jsonOut = true
			} else {
				dry = true
			}
			continue
		}
		if seen[f] && f != "--from-sql" && f != "--assume" {
			return invalid("duplicate " + f)
		}
		switch f {
		case "--workspace":
			seen[f] = true
			v, e := one(rest, &i, f)
			if e != nil {
				return nil, e
			}
			ws = v
		case "--resume":
			seen[f] = true
			v, e := one(rest, &i, f)
			if e != nil {
				return nil, e
			}
			resume = v
		case "--from-sql":
			v, e := one(rest, &i, f)
			if e != nil {
				return nil, e
			}
			if len(sql) >= 32 {
				return invalid("too many --from-sql files")
			}
			sql = append(sql, v)
		case "--assume":
			v, e := one(rest, &i, f)
			if e != nil {
				return nil, e
			}
			if len(assume) >= maxAssumptions || len([]byte(v)) > maxAssumptionBytes {
				return invalid("--assume exceeds limit")
			}
			if !regexp.MustCompile(`^[A-Za-z_][A-Za-z0-9_]*(?:\.[A-Za-z_][A-Za-z0-9_]*)?=[A-Za-z_][A-Za-z0-9_]*(?:\.[A-Za-z_][A-Za-z0-9_]*)?$`).MatchString(v) {
				return invalid("--assume must be column=table")
			}
			assume = append(assume, v)
		case "--output":
			seen[f] = true
			v, e := one(rest, &i, f)
			if e != nil {
				return nil, e
			}
			output = v
		case "--annotations":
			seen[f] = true
			v, e := one(rest, &i, f)
			if e != nil {
				return nil, e
			}
			annotations = v
		case "--reviewed-candidates":
			seen[f] = true
			v, e := one(rest, &i, f)
			if e != nil {
				return nil, e
			}
			reviewed = v
		}
	}
	if err := requireWorkspace(ws); err != nil {
		return nil, err
	}
	if resume != "" {
		if err := requireRun(resume); err != nil {
			return nil, err
		}
	}
	if action == "inspect" && resume != "" {
		return invalid("inspect does not accept --resume")
	}
	if action == "schema check" {
		if resume == "" {
			return invalid("schema check requires --resume")
		}
		if (annotations == "") != (reviewed == "") {
			return invalid("--annotations and --reviewed-candidates must be supplied together")
		}
		if reviewed != "" && !digestPattern.MatchString(reviewed) {
			return invalid("--reviewed-candidates must be sha256:<hex>")
		}
	}
	if action != "schema check" && (annotations != "" || reviewed != "") {
		return invalid("annotation options are valid only for schema check")
	}
	switch action {
	case "inspect":
		return InspectCommand{ws, jsonOut}, nil
	case "preprocess dwh":
		return DwhRequest{ws, resume, jsonOut}, nil
	case "schema suggest-fks":
		return SuggestFksRequest{ws, sql, assume, output, jsonOut}, nil
	case "schema check":
		return CheckSchemaRequest{ws, resume, annotations, reviewed, jsonOut}, nil
	case "index-schema":
		return IndexSchemaRequest{ws, jsonOut}, nil
	case "preprocess evidence":
		return EvidenceRequest{ws, dry, resume, jsonOut}, nil
	default:
		return RunRequest{ws, resume, jsonOut}, nil
	}
}
func min(a, b int) int {
	if a < b {
		return a
	}
	return b
}
func operationName(c Command) string {
	switch c.(type) {
	case InspectCommand:
		return "inspect"
	case DwhRequest:
		return "dwh"
	case SuggestFksRequest:
		return "suggest-fks"
	case CheckSchemaRequest:
		return "check"
	case IndexSchemaRequest:
		return "index-schema"
	case EvidenceRequest:
		return "evidence"
	default:
		return "run"
	}
}

func makeInput(c Command) (inputEnvelope, string, error) {
	env := inputEnvelope{SchemaVersion: 1, Operation: operationName(c)}
	var out string
	switch x := c.(type) {
	case InspectCommand:
		env.WorkspaceID = x.WorkspaceID
	case DwhRequest:
		env.WorkspaceID, env.Resume = x.WorkspaceID, x.Resume
	case SuggestFksRequest:
		env.WorkspaceID, env.Assume, out = x.WorkspaceID, x.Assume, x.Output
		if out != "" {
			if err := safeio.ValidateCanonicalOutputPath(out); err != nil {
				return env, "", errors.New("unsafe output file")
			}
		}
		var total int64
		for _, path := range x.FromSQL {
			b, e := safeio.ReadCanonicalUTF8(path, maxSQLFile)
			if e != nil {
				return env, "", errors.New("unsafe SQL input")
			}
			total += int64(len(b))
			if total > maxSQLTotal {
				return env, "", errors.New("SQL input exceeds limit")
			}
			base := path[strings.LastIndexAny(path, "/\\")+1:]
			env.SQL = append(env.SQL, sqlInput{base, base64.StdEncoding.EncodeToString(b), DigestBytes(b)})
		}
	case CheckSchemaRequest:
		env.WorkspaceID, env.Resume, env.ReviewedCandidates = x.WorkspaceID, x.Resume, x.ReviewedCandidates
		if x.Annotations != "" {
			b, e := safeio.ReadCanonicalUTF8(x.Annotations, maxAnnotations)
			if e != nil {
				return env, "", errors.New("unsafe annotation input")
			}
			base := x.Annotations[strings.LastIndexAny(x.Annotations, "/\\")+1:]
			env.Annotations = &annotationInput{base, base64.StdEncoding.EncodeToString(b), DigestBytes(b)}
		}
	case IndexSchemaRequest:
		env.WorkspaceID = x.WorkspaceID
	case EvidenceRequest:
		env.WorkspaceID, env.Resume, env.DryRun = x.WorkspaceID, x.Resume, x.DryRun
	case RunRequest:
		env.WorkspaceID, env.Resume = x.WorkspaceID, x.Resume
	default:
		return env, "", errors.New("unsupported workspace command")
	}
	payload, e := json.Marshal(env)
	if e != nil {
		return env, "", e
	}
	if len(payload) > maxRequest {
		return env, "", errors.New("request exceeds limit")
	}
	return env, string(payload), nil
}

var commonIngress = map[string]bool{"schemaVersion": true, "operation": true, "workspaceId": true}

func validateIngress(payload []byte, expected inputEnvelope) error {
	var raw map[string]json.RawMessage
	d := json.NewDecoder(bytes.NewReader(payload))
	d.UseNumber()
	if d.Decode(&raw) != nil || hasDuplicateJSONFields(payload) {
		return errors.New("invalid workspace request")
	}
	var extra any
	if d.Decode(&extra) != io.EOF {
		return errors.New("invalid workspace request")
	}
	var supplied inputEnvelope
	dec := json.NewDecoder(bytes.NewReader(payload))
	dec.DisallowUnknownFields()
	if dec.Decode(&supplied) != nil || supplied.SchemaVersion != 1 || supplied.Operation != expected.Operation || supplied.WorkspaceID != expected.WorkspaceID {
		return errors.New("invalid workspace request")
	}
	// Presence is part of the command-derived envelope contract. Decoding into Go values
	// alone would make an omitted field indistinguishable from an explicit zero/null value.
	var generatedRaw map[string]json.RawMessage
	if json.Unmarshal(mustJSON(expected), &generatedRaw) != nil || !sameJSONFieldSet(raw, generatedRaw) {
		return errors.New("workspace request does not match command")
	}
	if !reflect.DeepEqual(supplied, expected) {
		return errors.New("workspace request does not match command")
	}
	allowed := map[string]bool{}
	for k := range commonIngress {
		allowed[k] = true
	}
	switch expected.Operation {
	case "dwh", "run":
		if expected.Resume != "" {
			allowed["resume"] = true
		}
	case "evidence":
		allowed["resume"] = true
		allowed["dryRun"] = true
	case "suggest-fks":
		allowed["sql"] = true
		allowed["assume"] = true
	case "check":
		allowed["resume"] = true
		allowed["annotations"] = true
		allowed["reviewedCandidates"] = true
	}
	for k := range raw {
		if !allowed[k] {
			return errors.New("invalid workspace request")
		}
	}
	return nil
}

func Run(ctx context.Context, installation config.Installation, runner compose.Runner, command Command, stdin io.Reader) (Result, error) {
	env, generated, e := makeInput(command)
	if e != nil {
		return Result{}, e
	}
	payload := []byte(generated)
	if stdin != nil {
		payload, e = io.ReadAll(io.LimitReader(stdin, maxRequest+1))
		if e != nil || len(payload) > maxRequest {
			return Result{}, errors.New("request exceeds limit")
		}
		if e = validateIngress(payload, env); e != nil {
			return Result{}, e
		}
	}
	args := installation.ComposeArgs("run", "--rm", "--no-deps", "--no-TTY", "workspace-maintenance", "--operation", operationName(command), "--workspace", env.WorkspaceID)
	cr, runErr := runner.RunBounded(ctx, args, bytes.NewReader(payload), compose.CaptureLimits{StdoutBytes: maxResult, StderrBytes: 64 << 10})
	if errors.Is(runErr, compose.ErrOutputLimit) {
		return Result{}, runErr
	}
	if len(cr.Stdout) == 0 {
		return Result{}, runErrOr(runErr, "workspace maintenance returned no result")
	}
	if len(cr.Stdout) > maxResult {
		return Result{}, errors.New("invalid workspace result")
	}
	var raw map[string]json.RawMessage
	d := json.NewDecoder(strings.NewReader(cr.Stdout))
	if d.Decode(&raw) != nil || hasDuplicateJSONFields([]byte(cr.Stdout)) {
		return Result{}, errors.New("invalid workspace result")
	}
	var extra any
	if d.Decode(&extra) != io.EOF {
		return Result{}, errors.New("invalid workspace result")
	}
	var export *hostExport
	hasExport := false
	if h, ok := raw["hostExport"]; ok {
		hasExport = true
		if operationName(command) != "suggest-fks" {
			return Result{}, errors.New("invalid workspace result")
		}
		var candidate hostExport
		if hasDuplicateJSONFields(h) {
			return Result{}, errors.New("invalid host export")
		}
		exportDecoder := json.NewDecoder(bytes.NewReader(h))
		exportDecoder.DisallowUnknownFields()
		if exportDecoder.Decode(&candidate) != nil {
			return Result{}, errors.New("invalid host export")
		}
		var trailing any
		if exportDecoder.Decode(&trailing) != io.EOF {
			return Result{}, errors.New("invalid host export")
		}
		export = &candidate
		delete(raw, "hostExport")
	}
	b, _ := json.Marshal(raw)
	var result Result
	dec := json.NewDecoder(bytes.NewReader(b))
	dec.DisallowUnknownFields()
	if dec.Decode(&result) != nil {
		return Result{}, errors.New("invalid workspace result")
	}
	if e = validateResult(result, env.WorkspaceID, operationName(command)); e != nil {
		return Result{}, e
	}
	if !resultExitMatches(result.Status, cr.ExitCode) {
		return Result{}, runErrOr(runErr, "workspace result status does not match child exit")
	}
	if runErr != nil {
		// Exit 1/3 are represented by the validated public result; callers still
		// render it and select the corresponding process exit code.
		if result.Status != "blocked" && result.Status != "failed" {
			return Result{}, runErr
		}
	}
	if hasExport {
		if export == nil {
			return Result{}, errors.New("invalid candidate export")
		}
		if !candidateBoundToResult(*export, result) {
			return Result{}, errors.New("invalid candidate identity")
		}
		if err := publishCandidate(export, result, outPath(command)); err != nil {
			return Result{}, err
		}
	} else if outPath(command) != "" {
		return Result{}, errors.New("candidate export is required")
	}
	return result, nil
}
func runErrOr(e error, msg string) error {
	if e != nil {
		return e
	}
	return errors.New(msg)
}
func outPath(c Command) string {
	if x, ok := c.(SuggestFksRequest); ok {
		return x.Output
	}
	return ""
}
func candidateBoundToResult(x hostExport, result Result) bool {
	if result.RunID == "" {
		return false
	}
	for _, artifact := range result.ArtifactIdentities {
		if artifact.Kind == "fk-candidates" && artifact.Digest == x.SHA256 {
			return true
		}
	}
	return false
}

func publishCandidate(x *hostExport, result Result, path string) error {
	if x.MediaType != "application/yaml" && x.MediaType != "text/yaml" {
		return errors.New("invalid candidate export")
	}
	if !digestPattern.MatchString(x.SHA256) {
		return errors.New("invalid candidate export")
	}
	b, e := base64.StdEncoding.DecodeString(x.ContentBase64)
	if e != nil || len(b) > maxCandidate || !utf8.Valid(b) {
		return errors.New("invalid candidate export")
	}
	if DigestBytes(b) != x.SHA256 {
		return errors.New("invalid candidate export")
	}
	var doc any
	if yaml.Unmarshal(b, &doc) != nil {
		return errors.New("invalid candidate export")
	}
	if result.RunID == "" || !runIDPattern.MatchString(result.RunID) {
		return errors.New("invalid candidate identity")
	}
	if path == "" {
		return nil
	}
	if e = safeio.WriteCanonicalExclusive(path, b, 0o600); e != nil {
		return errors.New("unsafe output file")
	}
	return nil
}
func validateResult(r Result, workspace, operation string) error {
	if r.SchemaVersion != 1 || r.WorkspaceID != workspace || r.Operation != operation || !validStatus(r.Status) || !validCode(r.Code) || !revisionPattern.MatchString(r.WorkspaceRevision) || !revisionPattern.MatchString(r.DescriptorBlob) {
		return errors.New("invalid workspace result")
	}
	if r.RunID != "" && !runIDPattern.MatchString(r.RunID) {
		return errors.New("invalid workspace result")
	}
	for k, v := range r.ChildRuns {
		if k == "" || !runIDPattern.MatchString(v) {
			return errors.New("invalid workspace result")
		}
	}
	for _, a := range r.ArtifactIdentities {
		if a.Kind == "" || !digestPattern.MatchString(a.Digest) {
			return errors.New("invalid workspace result")
		}
	}
	if r.CompletedStages == nil {
		return errors.New("invalid workspace result")
	}
	if r.Status == "blocked" && (r.Code != "manual_review_required" && r.Code != "evidence_materialization_required" && r.Code != "preprocessing_conflict" && r.Code != "preprocessing_resume_mismatch" && r.Code != CodeRegistryBootstrapRecoveryConflict) {
		return errors.New("invalid workspace result")
	}
	if r.Status != "blocked" && r.Code == CodeRegistryBootstrapRecoveryConflict {
		return errors.New("invalid workspace result")
	}
	if (r.Status == "succeeded" || r.Status == "unchanged" || r.Status == "dry_run") && r.Code != "ok" {
		return errors.New("invalid workspace result")
	}
	if r.Status == "failed" && (r.Code == "ok" || r.Code == "manual_review_required" || r.Code == "evidence_materialization_required" || r.Code == "preprocessing_conflict" || r.Code == "preprocessing_resume_mismatch") {
		return errors.New("invalid workspace result")
	}
	return nil
}

func resultExitMatches(status string, exitCode int) bool {
	switch status {
	case "succeeded", "unchanged", "dry_run":
		return exitCode == 0
	case "blocked":
		return exitCode == 3
	case "failed":
		return exitCode == 1
	default:
		return false
	}
}

func hasDuplicateJSONFields(payload []byte) bool {
	decoder := json.NewDecoder(bytes.NewReader(payload))
	decoder.UseNumber()
	if !scanJSONValue(decoder) {
		return true
	}
	_, err := decoder.Token()
	return err != io.EOF
}

func scanJSONValue(decoder *json.Decoder) bool {
	token, err := decoder.Token()
	if err != nil {
		return false
	}
	delimiter, isDelimiter := token.(json.Delim)
	if !isDelimiter {
		return true
	}
	switch delimiter {
	case '{':
		seen := map[string]bool{}
		for decoder.More() {
			keyToken, keyErr := decoder.Token()
			key, ok := keyToken.(string)
			if keyErr != nil || !ok || seen[key] || !scanJSONValue(decoder) {
				return false
			}
			seen[key] = true
		}
		_, err = decoder.Token()
		return err == nil
	case '[':
		for decoder.More() {
			if !scanJSONValue(decoder) {
				return false
			}
		}
		_, err = decoder.Token()
		return err == nil
	default:
		return false
	}
}

func sameJSONFieldSet(a, b map[string]json.RawMessage) bool {
	if len(a) != len(b) {
		return false
	}
	for key := range a {
		if _, ok := b[key]; !ok {
			return false
		}
	}
	return true
}

func mustJSON(v any) []byte {
	b, _ := json.Marshal(v)
	return b
}
func DigestBytes(b []byte) string { s := sha256.Sum256(b); return "sha256:" + hex.EncodeToString(s[:]) }
func validStatus(v string) bool {
	switch v {
	case "succeeded", "unchanged", "dry_run", "blocked", "failed":
		return true
	}
	return false
}
func validCode(v string) bool {
	switch v {
	case "ok", "workspace_not_found", "workspace_not_activatable", "binding_missing", "preprocessing_conflict", "preprocessing_resume_mismatch", "manual_review_required", "evidence_materialization_required", "effective_config_mismatch", "semantic_index_incompatible", "annotation_invalid", "egress_policy_refused", CodeRegistryBootstrapRecoveryConflict:
		return true
	}
	return false
}